Implement a Network in Network CNN model using pytorch-lightning

[sushmit86]

I am trying to implement a NiN model. Basically trying to replicate code from d2lbook Here is my code.

import pandas as pd
import torch
from torch import nn
import torchmetrics
from torchvision import transforms
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as pl
from torchvision.datasets import FashionMNIST
import wandb
from pytorch_lightning.loggers import WandbLogger
wandb.login()

## class definition
class Lightning_nin(pl.LightningModule):
  def __init__(self):
    super().__init__()
    self.accuracy = torchmetrics.Accuracy(top_k=1)
    self.model = nn.Sequential(
                self.nin_block(1, 96, kernel_size=11, strides=4, padding=0),
                nn.MaxPool2d(3, stride=2),
                self.nin_block(96, 256, kernel_size=5, strides=1, padding=2),
                nn.MaxPool2d(3, stride=2),
                self.nin_block(256, 384, kernel_size=3, strides=1, padding=1),
                nn.MaxPool2d(3, stride=2), nn.Dropout(0.5),
                # There are 10 label classes
                self.nin_block(384, 10, kernel_size=3, strides=1, padding=1),
                nn.AdaptiveAvgPool2d((1, 1)),
                # Transform the four-dimensional output into two-dimensional output with a
                # shape of (batch size, 10)
                nn.Flatten())
    for layer in self.model:
      if type(layer) == nn.Linear or type(layer) == nn.Conv2d:
        nn.init.xavier_uniform_(layer.weight)  
  def nin_block(self,in_channels, out_channels, kernel_size, strides, padding):
      return nn.Sequential(
          nn.Conv2d(in_channels, out_channels, kernel_size, strides, padding),
          nn.ReLU(), nn.Conv2d(out_channels, out_channels, kernel_size=1),
          nn.ReLU(), nn.Conv2d(out_channels, out_channels, kernel_size=1),
          nn.ReLU())
      
  def forward(self, x):
      x = self.model(x)
      return x
 
  def loss_fn(self,logits,y):
    loss = nn.CrossEntropyLoss()
    return loss(logits,y)
    
  def training_step(self,train_batch,batch_idx):
    X, y = train_batch
    logits = self.forward(X)
    loss = self.loss_fn(logits,y)
    self.log('train_loss',loss)
    m = nn.Softmax(dim=1)
    output = m(logits)
    self.log('train_acc',self.accuracy(output,y))
    return loss
  
  def validation_step(self,val_batch,batch_idx):
    X,y = val_batch
    logits = self.forward(X)
    loss = self.loss_fn(logits,y)
    self.log('test_loss',loss)
    m = nn.Softmax(dim=1)
    output = m(logits)
    self.log('test_acc',self.accuracy(output,y))
  
  def configure_optimizers(self):
    optimizer = torch.optim.SGD(self.model.parameters(),lr= 0.1)
    return optimizer
  
class Light_DataModule(pl.LightningDataModule):
  def __init__(self,resize= None):
    super().__init__()
    if resize:
      self.resize = resize
 
  def setup(self, stage):
    # transforms for images
    trans = [transforms.ToTensor()]
    if self.resize:
      trans.insert(0, transforms.Resize(self.resize))
    trans = transforms.Compose(trans)
    # prepare transforms standard to MNIST
    self.mnist_train = FashionMNIST(root="../data", train=True, download=True, transform=trans)
    self.mnist_test = FashionMNIST(root="../data", train=False, download=True, transform=trans)
 
  def train_dataloader(self):
    return DataLoader(self.mnist_train, batch_size=128,shuffle=True,num_workers=4)
 
  def val_dataloader(self):
    return DataLoader(self.mnist_test, batch_size=128,num_workers=4)

## Train model
data_module = Light_DataModule(resize=224)
wandb_logger = WandbLogger(project="d2l",name ='NIN')
model  = Lightning_nin()
trainer = pl.Trainer(logger=wandb_logger,max_epochs=4,gpus=1,progress_bar_refresh_rate =1)
trainer.fit(model, data_module)
wandb.finish()

After running the code I am only getting an accuracy of 0.1. Not sure where I am going wrong. I have been able to implement other CNN (like VGG) using same template. Not sure where I am going wromg.

[sushmit86]

@williamFalcon Reaching out to you for help on this

[awaelchli]

Do I see correctly, you want logits from your forward() but you won't get them, because the last nin block has a relu and the only layers that follows is the avg pool and the flatten.

You posted the link to the architecture, they even have code there, so you have to read it carefully and pay attention to how the architecture is built. Check that the inputs range and output range makes sense.

[sushmit86]

Thanks for your reply. But I am using nn.CrossEntropyLoss() as the loss so it should take care of the logit. Also, the architecture is exactly the same as the posted link, I also checked the shape at each layer. Here is the output

Sequential output shape:     torch.Size([1, 96, 54, 54])
MaxPool2d output shape:      torch.Size([1, 96, 26, 26])
Sequential output shape:     torch.Size([1, 256, 26, 26])
MaxPool2d output shape:      torch.Size([1, 256, 12, 12])
Sequential output shape:     torch.Size([1, 384, 12, 12])
MaxPool2d output shape:      torch.Size([1, 384, 5, 5])
Dropout output shape:        torch.Size([1, 384, 5, 5])
Sequential output shape:     torch.Size([1, 10, 5, 5])
AdaptiveAvgPool2d output shape:      torch.Size([1, 10, 1, 1])
Flatten output shape:        torch.Size([1, 10])