Is my converted TensorFlow model to PyTorch Correct?

mustansarsaeed: <p>Hi,<br>
I am trying to convert the <code>TensorFlow</code> model to PyTorch. Can anyone please comment that both models are correct or have any issues because output is different. My pytorch model is giving constant loss.</p>
<p><strong>TensorFlow</strong></p>
<pre><code class="lang-auto">features = tf.placeholder(
                tf.float32, shape=[None, IMAGE_SIZE * IMAGE_SIZE], name='features')
labels = tf.placeholder(tf.int64, shape=[None], name='labels')
input_layer = tf.reshape(features, [-1, IMAGE_SIZE, IMAGE_SIZE, 1])
conv1 = tf.layers.conv2d(
              inputs=input_layer,
              filters=32,
              kernel_size=[5, 5],
              padding="same",
              activation=tf.nn.relu )
pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
conv2 = tf.layers.conv2d(
                inputs=pool1,
                filters=64,
                kernel_size=[5, 5],
                padding="same",
                activation=tf.nn.relu )
pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
dense = tf.layers.dense(inputs=pool2_flat, units=2048, name='dense1', 
             kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001), 
             bias_regularizer=tf.contrib.layers.l2_regularizer(0.001))
act_1 = tf.nn.relu(dense)
logits = tf.layers.dense(inputs=act_1, units=self.num_classes)
predictions = {
              "classes": tf.argmax(input=logits, axis=1),
              "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
            }
loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
            # TODO: Confirm that opt initialized once is ok?
train_op = self.optimizer.minimize(
                loss=loss,
                global_step=tf.train.get_global_step())
eval_metric_ops = tf.count_nonzero(tf.equal(labels, predictions["classes"]))
</code></pre>
<p><strong>Equivalent PyTorch</strong></p>
<pre><code class="lang-auto">class FemnistNet(nn.Module):
    def __init__(self):
        super(FemnistNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2)
        self.pool1 = nn.MaxPool2d(2, stride=2, )
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2)

        self.pool2 = nn.MaxPool2d(2, stride=2)
        self.fc1 = nn.Linear(3136, 2048)
        self.fc2 = nn.Linear(2048 ,62)

        
    def forward(self, x):
        x = x.view(-1, 1, 28, 28)
        x = self.conv1(x)
        x = th.nn.functional.relu(x)

        x = self.pool1(x)

        x=self.conv2(x)
        x = th.nn.functional.relu(x)
        
        x = self.pool2(x)
        
        x = x.flatten(start_dim=1)
        
        x = self.fc1(x)
        l1_activations = th.nn.functional.relu(x)
        
        x = self.fc2(l1_activations)

        x = x.softmax()

        return x, l1_activations
</code></pre>
<blockquote>
<p>My training accuracy and loss are constant in the PyTorch model. What can be the issue?</p>
</blockquote>

ptrblck: <p>I don’t know what your use case and in particular loss function is, but the last <code>softmax</code> operation is often wrong. In case you are working on a multi-class classification use case using <code>nn.CrossEntropyLoss</code>, note that this loss function expects raw logits so you would have to remove the <code>x = x.softmax()</code> call.<br>
Also, make sure to add the weight decay to the <code>optimizer</code> for the same parameters as was done in TF.</p>