Moving TF code to PyTorch

Gabor_Erdos: <p>Hi! I am completely new to PyTorch, I would like to move my TF code to PyTorch, and I think I am missing something.</p>
<p>I have X as input and Y as output. X is a time series data, on which I would like to do 1D convolution. Y is just a plain number.</p>
<p>X has a shape of (1050589, 81, 21). I have 1050589 experiments, each experiment has 81 timestamps and each timestamp has 21 points of data. This is the required format for TF, but as far as I was able to get out in PyTorch the time dimension should be the last one.</p>
<p>I have my data in a numpy array, so first I transformed the data to fit PyTorch, and also transformed into a list.</p>
<pre><code class="lang-auto">a = []
for n, i in enumerate(X):
    a.append([X[n].T, Y[n]])

train_data = DataLoader(a, batch_size=128)
</code></pre>
<p>My model looks like this:</p>
<pre><code class="lang-auto">class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.linear_relu_stack = nn.Sequential(
            nn.Conv1d(EMBED_SIZE, 32, 7, padding='same'),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(81*32, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
        )

    def forward(self, x):
        logits = self.linear_relu_stack(x)
        return logits.double()
</code></pre>
<p>The architecture is simple, as I want to keep it the same as I have in Tensorflow. One convolution with a kernel of 7 and 32 channels, followed by a dense layer and a single output layer.</p>
<p>Same network in Tensorflow:</p>
<pre><code class="lang-auto">def conv_1d_model():
    model = Sequential(name="model_conv1D")
    model.add(Conv1D(filters=32, kernel_size=7, activation='relu', input_shape=(81, 21), padding="same"))
    model.add(Flatten())
    model.add(Dense(32, activation='relu'))
    model.add(Dense(1))
    return model
</code></pre>
<p>Now when I try to optimize this network in PyTorch my losses are all over the place, not decreasing at all, while in TensorFlow it runs perfectly well.</p>
<p>I am sure I am missing something, can anyone point me in the right direction?</p>
<p>My optimization function in PyTorch:</p>
<pre><code class="lang-auto">model = NeuralNetwork()

loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)


def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred = torch.squeeze(model(X))  # I was getting a warning about the pred being in different shape than y, so I squeezed it
        loss = loss_fn(pred, y)
        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 10 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:&gt;7f}  [{current:&gt;5d}/{size:&gt;5d}]")
</code></pre>
<p>Optimization in Tensorflow</p>
<pre><code class="lang-auto">model = conv_1d_model()
opt = Adam(learning_rate=learning_rate)
model.compile(loss='mse', optimizer=opt, metrics=['mae'])

model_history = model.fit(X, Y, validation_split=0.2, epochs=epochs, batch_size=batch_size, verbose=1)
</code></pre>

nivek: <p>From a first glance, it looks like your samples are not being shuffled. Is there a reason for that?</p>

Gabor_Erdos: <p>I dont think it would matter, these are actually protein sequences embedded, also I think the TF code lacks shuffling as well, so that should not yield these strange results.</p>

ptrblck: <pre><code class="lang-auto"># I was getting a warning about the pred being in different shape than y, so I squeezed it
</code></pre>
<p>Could you check the shape of the model output and the targets?<br>
I would assume <code>nn.MSELoss</code> was warning about a potentially unwanted broadcasting, so could you calculate the loss for one batch manually and compare it to the output of <code>MSELoss</code>?</p>

Gabor_Erdos: <p>While the optimization is running the shape of y is [256], while the returned shape of pred was [[256, 1]]. After I squeezed, the warning was gone.</p>
<p>How can I check the loss on one batch manually?</p>

ptrblck: <p>You could write the loss function e.g. via <code>((x - y)**2).mean()</code> and check if the intermediates would be broadcasted. However, since you’ve squeezed one tensor, I would assume the shapes are equal (you could <code>unsqueeze</code> the other tensor in <code>dim1</code> which should also work and might be more explicit).<br>
Assuming the loss calculation is correct I would recommend to try to overfit a small dataset (e.g. just 10 samples) and see if your model is able to do so.</p>

Gabor_Erdos: <p>Dear ptrblck!</p>
<p>Thanks a lot for helping me! I tried what you said, and to my surprise with only 10 elements in X and Y the model converges perfectly well, the loss goes to 1e-2 in 100 epochs.</p>
<p>I thought it would fail. Now have have absolutely no idea what causes the difference between the TF run and the PyTorch one.</p>

ptrblck: <p>Are you seeing the same final loss in TF as ~1e-2 for these 10 samples as well?<br>
If so, could you try to scale up the training use case in PyTorch again and see when it breaks?<br>
If the initial training looks good, I would start checking the parameter initialization next.</p>