<jupyter_start><jupyter_text>Using Yuzu with Your ModelA challenge in using Yuzu is that replacing operations within a graph is somewhat difficult to implement efficiently: not only do the convolution and max-pool operations need to be rewritten, but the precomputed metadata needs to be passed into layers as appropriate. Accordingly, Yuzu's implementation ignores the forward pass in the written model and instead iterates over the layers in the `model.children()` generator, applying or replacing them as necessary.This notebook will show how to use Yuzu in practice. First, we will show some of the limitations that Yuzu has. Then, we will show how to overcome many of these limitations easily.<jupyter_code>%pylab inline
import seaborn; seaborn.set_style('whitegrid')
import torch<jupyter_output>Populating the interactive namespace from numpy and matplotlib<jupyter_text>Limitations of YuzuUnfortunately, iterating over `model.children()` is conceptually easy but requires that the models are sequential. A major consequence of this is that no multi-input or multi-output networks are supported right now (unless the inputs or outputs can be represented in a single tensor). Also, no networks with multiple paths are supported right now, including residual connections. A second type of limitation is that no operations can be performed in the forward pass other than iterating through the layers in a sequential manner. This means no flattening, reshaping, cropping, adding operations, referring to the same layer multiple times, etc. Any manipulation of the data must occur within the context of layers.Let's start off with a minimal implementation of a small BPNet model, implemented as one might without considering these constraints.<jupyter_code>import torch

class BPNet(torch.nn.Module):
    def __init__(self, n_filters=64, n_layers=3):
        super(BPNet, self).__init__()
        self.trimming = 2 ** n_layers + 3
        self.n_filters = n_filters
        self.n_layers = n_layers

        self.iconv = torch.nn.Conv1d(4, n_filters, kernel_size=7, padding=3)
        self.rconvs = torch.nn.ModuleList([
            torch.nn.Conv1d(n_filters, n_filters, kernel_size=3, padding=2**i, dilation=2**i) for i in range(1, self.n_layers+1)
        ])

        self.fconv = torch.nn.Conv1d(n_filters, 1, kernel_size=7, padding=3)
        self.relu = torch.nn.ReLU()
        self.logsoftmax = torch.nn.LogSoftmax(dim=-1)

    def forward(self, X):
        start, end = self.trimming, X.shape[2] - self.trimming
        X = self.relu(self.iconv(X))
        for i in range(self.n_layers):
            X_conv = self.relu(self.rconvs[i](X))
            X = torch.add(X, X_conv)

        X = self.fconv(X)
        X = X[:, :, start:end]
        y_profile = self.logsoftmax(X) 
        return y_profile<jupyter_output><empty_output><jupyter_text>Unfortunately, using Yuzu requires that a few things need to be changed in this implementation. The first is that ModuleList layer needs to be flattened out. The second is that each ReLU has to be explicitly written rather than having a single ReLU object referenced multiple times. Next, the cropping operation has to be moved into an explicit layer that performs cropping and is specified in the `__init__` method. Finally, and most unfortunately, residual operations are not currently supported by Yuzu.Here is what a model that would work for Yuzu would look like:<jupyter_code>class Crop(torch.nn.Module):
    def __init__(self, trimming):
        super(Crop, self).__init__()
        self.trimming = trimming
    
    def forward(self, X):
        start, end = self.trimming, X.shape[2] - self.trimming
        return X[:, :, start:end]

class BPNet2(torch.nn.Module):
    def __init__(self, n_filters=64, n_layers=3):
        super(BPNet2, self).__init__()
        self.trimming = 2 ** n_layers + 3
        self.n_filters = n_filters
        self.n_layers = n_layers

        self.iconv = torch.nn.Conv1d(4, n_filters, kernel_size=7, padding=3)
        self.relu1 = torch.nn.ReLU()
        
        self.rconv1 = torch.nn.Conv1d(n_filters, n_filters, kernel_size=3, padding=2, dilation=2)
        self.relu2 = torch.nn.ReLU()
        
        self.rconv2 = torch.nn.Conv1d(n_filters, n_filters, kernel_size=3, padding=4, dilation=4)
        self.relu3 = torch.nn.ReLU()
        
        self.rconv3 = torch.nn.Conv1d(n_filters, n_filters, kernel_size=3, padding=8, dilation=8)
        self.relu4 = torch.nn.ReLU()
        
        self.fconv = torch.nn.Conv1d(n_filters, 1, kernel_size=7, padding=3)
        self.crop = Crop(self.trimming)
        self.logsoftmax = torch.nn.LogSoftmax(dim=-1)

    def forward(self, X):
        X = self.relu1(self.iconv(X))
        X = self.relu2(self.rconv1(X))
        X = self.relu3(self.rconv2(X))
        X = self.relu4(self.rconv3(X))
        X = self.fconv(X)
        X = self.crop(X)
        X = self.logsoftmax(X)
        return X
    
model = BPNet2()<jupyter_output><empty_output><jupyter_text>Given this model, we can run through the Yuzu pipeline just as before. All we have to do is make sure that the custom layers are passed in as terminal layers if you need to stop usin the delta tensor representation internally and instead stop running the Yuzu procedure once the layer is encountered.<jupyter_code>from yuzu import yuzu_ism, precompute

seq_len = 500

precomputation = precompute(model, seq_len, alpha=1.05, terminal_layers=(Crop,))
idxs = numpy.random.RandomState(0).randn(4, seq_len).argmax(axis=0)
X = numpy.zeros((1, 4, seq_len), dtype='float32')
X[0, idxs, numpy.arange(seq_len)] = 1

%timeit yuzu_ism(model, X, precomputation, terminal_layers=(Crop,))<jupyter_output>861 ms ± 202 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>And we can compare the time to the naive ISM method.<jupyter_code>from yuzu.naive_ism import naive_ism

%timeit naive_ism(model, X)<jupyter_output>3.19 s ± 215 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>And we can get the scores as normal. Just make sure to pass in the custom layer as a terminal layer.<jupyter_code>yuzu_ism_scores = yuzu_ism(model, X, precomputation, terminal_layers=(Crop,))
naive_ism_scores = naive_ism(model, X)
plt.scatter(yuzu_ism_scores.flatten(), naive_ism_scores.flatten())
plt.plot([0, 0.016], [0, 0.016])
plt.title("ISM Scores Match Across Methods", fontsize=14)
plt.xlabel("Yuzu ISM Scores", fontsize=12)
plt.ylabel("Naive ISM Scores", fontsize=12)
plt.show()<jupyter_output><empty_output><jupyter_text>Finally, we can confirm that we get the same ISM scores using Yuzu and the naive method on this model.<jupyter_code>numpy.abs(naive_ism_scores - yuzu_ism_scores).mean()<jupyter_output><empty_output><jupyter_text>Using Yuzu with a Pre-trained ModelNow, let's say that you've already trained your model and you don't want to have to rewrite and retrain your model just to use Yuzu. Fortunately, you can usually transfer the weights of the learned parameters over to a new model definition that includes the changes that you need.For instance, let's say you've trained a DeepSEA model with the following implementation:<jupyter_code>class DeepSEA(torch.nn.Module):
    def __init__(self, random_state=0):
        super(DeepSEA, self).__init__()
        torch.manual_seed(random_state)

        self.conv1 = torch.nn.Conv1d(4, 32, kernel_size=9, padding=4)
        self.relu1 = torch.nn.ReLU()
        self.mp1 = torch.nn.MaxPool1d(4)

        self.conv2 = torch.nn.Conv1d(32, 48, kernel_size=9, padding=4)
        self.relu2 = torch.nn.ReLU()
        self.mp2 = torch.nn.MaxPool1d(4)

        self.conv3 = torch.nn.Conv1d(48, 96, kernel_size=9, padding=4)
        self.relu3 = torch.nn.ReLU()

        self.fc = torch.nn.Linear((250 // 4 // 4) * 96, 925)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, X):
        X = self.mp1(self.relu1(self.conv1(X)))
        X = self.mp2(self.relu2(self.conv2(X)))
        X = self.relu3(self.conv3(X))
        X = X.permute(0, 2, 1).contiguous().view(X.shape[0], -1)

        X = self.sigmoid(self.fc(X))
        X = X.unsqueeze(1)
        return X
    
model = DeepSEA()<jupyter_output><empty_output><jupyter_text>In order to get a model trained with this implementation to work with Yuzu, all we need to do is redefine the implementation with the fixes and transfer the weights over as appropriate. First, let's write out an implementation that would work with Yuzu using two built-in layers for common operations.<jupyter_code>from yuzu.models import Flatten, Unsqueeze

class DeepSEA2(torch.nn.Module):
    def __init__(self, random_state=0):
        super(DeepSEA2, self).__init__()
        torch.manual_seed(random_state)

        self.conv1 = torch.nn.Conv1d(4, 32, kernel_size=9, padding=4)
        self.relu1 = torch.nn.ReLU()
        self.mp1 = torch.nn.MaxPool1d(4)

        self.conv2 = torch.nn.Conv1d(32, 48, kernel_size=9, padding=4)
        self.relu2 = torch.nn.ReLU()
        self.mp2 = torch.nn.MaxPool1d(4)

        self.conv3 = torch.nn.Conv1d(48, 96, kernel_size=9, padding=4)
        self.relu3 = torch.nn.ReLU()

        self.flatten = Flatten()
        self.fc = torch.nn.Linear((250 // 4 // 4) * 96, 925)
        self.sigmoid = torch.nn.Sigmoid()
        self.unsqueeze = Unsqueeze(1)

    def forward(self, X):
        X = self.mp1(self.relu1(self.conv1(X)))
        X = self.mp2(self.relu2(self.conv2(X)))
        X = self.relu3(self.conv3(X))
        X = self.flatten(X)

        X = self.sigmoid(self.fc(X))
        X = self.unsqueeze(X)
        return X<jupyter_output><empty_output><jupyter_text>Now, let's create an instance of the new model and transfer the parameters over.<jupyter_code>model2 = DeepSEA2()
model2.load_state_dict(model.state_dict())<jupyter_output><empty_output><jupyter_text>We can demonstrate that the outputs from the model are the same:<jupyter_code>idxs = numpy.random.RandomState(0).randn(4, 250).argmax(axis=0)
X = numpy.zeros((1, 4, 250), dtype='float32')
X[0, idxs, numpy.arange(250)] = 1
X_torch = torch.tensor(X)

x = model(X_torch).detach().numpy()
y = model2(X_torch).detach().numpy()

numpy.abs(x - y).sum()<jupyter_output>/users/jmschr/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py:652: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  /pytorch/c10/core/TensorImpl.h:1156.)
  return torch.max_pool1d(input, kernel_size, stride, padding, dilation, ceil_mode)<jupyter_text>We can also show that the naive ISM outputs are the same for both models, and that the Yuzu ISM output matches the naive ISM one.<jupyter_code>nism1 = naive_ism(model, X)
nism2 = naive_ism(model2, X)

precomputation = precompute(model2, 250)
yism = yuzu_ism(model2, X, precomputation)

numpy.abs(nism1 - nism2).sum(), numpy.abs(yism - nism2).sum()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10409260778415824, 'fraction_numerical': 0.0245191869474068, 'mean_word_length': 3.721170395869191, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '441329', 'license': 'MIT', 'n_tokens_mistral': 3820, 'n_tokens_neox': 3628, 'n_words': 1073, 'path': 'tutorials/3. Using Yuzu with Your Model.ipynb', 'repo_name': 'kundajelab/yuzu', 'size': '10971'}