Completely different results using Tensorflow and Pytorch for MobilenetV3 Small

Kaustubh_Kulkarni: <p>I am using transfer learning from MobileNetV3 Small to predict 5 different points on an image. I am doing this as a regression task.</p>
<p>For both models:</p>
<ul>
<li>Setting the last 50 layers trainable and adding the same fully connected layers to the end.</li>
<li>Learning rate <code>3e-2</code></li>
<li>Batch size 32</li>
<li>Adam optimizer with the same betas</li>
<li>100 epochs</li>
<li>The inputs consist of RGB unscaled images</li>
</ul>
<h1><a name="pytorch-1" class="anchor" href="#pytorch-1"></a>Pytorch</h1>
<h2><a name="model-2" class="anchor" href="#model-2"></a>Model</h2>
<pre><code class="lang-auto">def _init_weights(m):
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)
        m.bias.data.fill_(0.01)

def get_mob_v3_small():
    model = torchvision.models.mobilenet_v3_small(pretrained=True)
    children_list = get_children(model)
          for c in children_list[:-50]:
                for p in c.parameters():
                    p.requires_grad = False

    return model

class TransferMobileNetV3_v2(nn.Module):
    def __init__(self,
            num_keypoints: int = 5):
        super(TransferMobileNetV3_v2, self).__init__()

        self.classifier_neurons = num_keypoints*2
        self.base_model = get_mob_v3_small()

        self.base_model.classifier = nn.Sequential(
                                            nn.Linear(in_features=1024, out_features=1024),
                                            nn.ReLU(),
                                            nn.Linear(in_features=1024, out_features=512),
                                            nn.ReLU(),
                                            nn.Linear(in_features=512, out_features=self.classifier_neurons)
                                        )

        self.base_model.apply(_init_weights)

    def forward(self, x):
        out = self.base_model(x)
        return out
</code></pre>
<h2><a name="training-script-3" class="anchor" href="#training-script-3"></a>Training Script</h2>
<pre><code class="lang-auto">def train(net, trainloader, testloader, train_loss_fn, optimizer, scaler, args):

    len_dataloader = len(trainloader)

    for epoch in range(1, args.epochs+1):
        net.train()

        for batch_idx, sample in enumerate(trainloader):

            inputs, labels = sample
            inputs, labels = inputs.to(args.device), labels.to(args.device)

            optimizer.zero_grad()

            with torch.cuda.amp.autocast(args.use_amp):
                prediction = net(inputs)
                loss = train_loss_fn(prediction, labels)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

def main():
    args = make_args_parser()
    args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    seed = args.seed
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)

    loss_fn = nn.MSELoss()
    optimizer = optim.Adam(net.parameters(), lr=3e-2,
                        betas=(0.9, 0.999))
    scaler = torch.cuda.amp.GradScaler(enabled=args.use_amp)

    train(net, train_loader, test_loader, loss_fn, optimizer, scaler, args)

</code></pre>
<h1><a name="tensorflow-4" class="anchor" href="#tensorflow-4"></a>Tensorflow</h1>
<h2><a name="model-5" class="anchor" href="#model-5"></a>Model</h2>
<pre><code class="lang-auto">base_model = tf.keras.applications.MobileNetV3Small(weights='imagenet',input_shape=(224,224,3))

x_in = base_model.layers[-6].output

x = Dense(units=1024, activation="relu")(x_in)
x = Dense(units=512, activation="relu")(x)
x = Dense(units=10, activation="linear")(x)

model = Model(inputs=base_model.input, outputs=x)

for layer in model.layers[:-50]:
    layer.trainable=False
</code></pre>
<h2><a name="training-script-6" class="anchor" href="#training-script-6"></a>Training Script</h2>
<pre><code class="lang-auto">model.compile(loss = "mse",
              optimizer = tf.keras.optimizers.Adam(learning_rate=3e-2))

history = model.fit(input_numpy, output_numpy,
                  verbose=1,
                  batch_size=32, epochs=100,validation_split = 0.2)
</code></pre>
<h1><a name="results-7" class="anchor" href="#results-7"></a>Results</h1>
<ul>
<li>The PyTorch model predicts one single point around the center for all 5 different points.</li>
<li>The Tensorflow model predicts the points quite well and are quite accurate.</li>
<li>The loss in the Pytorch model is much higher than the Tensorflow model.</li>
</ul>
<p>Please do let me know what is going wrong as I am trying my best to shift to PyTorch for this work and I need this model to give me similar/identical results.</p>
<blockquote>
<p>Note: I also noticed that the MobileNetV3 Small model seems to be different in PyTorch and different in Tensorflow. I do not know if am interpreting it wrong, but I’m putting it here just in case.</p>
</blockquote>