Model not converging

Keras to Pytorch

I am train a model that I trained with 93% accuracy on Keras.
The problem is my model is not converging and giving an accuracy of 20%

The code for my model is as follows:

class SmallClassifier(nn.Module):

def __init__(self,convs,class_size,input_shape=(1, 28, 28),batch_size=64):
    super(SmallClassifier,self).__init__()
    self.convs = convs
    self.batch_size = batch_size
    self.class_size = class_size
    self.conv2_drop = nn.Dropout2d()
    self.conv_output = self._get_conv_output(input_shape)
    self.classifier = nn.Sequential(
        nn.Linear(self.conv_output, 512),
        nn.Linear(512,class_size),
        nn.LogSoftmax()
    )

    self._initialize_weights()
    

''' Forward pass with Global Average Softmax'''
def forward(self,x):
    x = self.convs(x)
    x = x.view(-1,self.conv_output)
    x = self.classifier(x)
    return x

# generate input sample and forward to get shape
def _get_conv_output(self, shape):
    
    input = Variable(torch.rand(self.batch_size, *shape))
    output_feat = self._forward_features(input)
    n_size = output_feat.data.view(self.batch_size, -1).size(1)
    return n_size

def _forward_features(self, x):
    x = self.convs(x)
    x = self.conv2_drop(x)
    return x

''' Init Weights initialising by using w = np.random.randn(n) * sqrt(2.0/n) according to  He et al.. '''
def _initialize_weights(self):
     for m in self.modules():
        if isinstance(m, nn.Conv2d):
            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            m.weight.data.normal_(0, math.sqrt(2. / n))
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.fill_(1)
            m.bias.data.zero_()

`

The self.conv gets all the conv nets , which I took from the vgg.py template on pytorch

The code for getting self.conv is as follows:

cfg = {
    'yolo_full': [32,'M', 64, 'M', 128,64,128,'M', 256,128, 256,'M',512,256,512,256,512,'M',1024,512,1024,512,1024],
    'small_classifier':[32,64,64,'M']
}

filter_size = {
    'yolo_full': [3,'M',3,'M',3,1,3,'M',3,1,3,'M',3,1,3,1,3,'M',3,1,3,1,3],
    'small_classifier':[3,3,3,'M']
}



def make_layers(cfg,filter_size, batch_norm=True,in_channels=3):
    layers = []

    i  = 0
    for v in cfg:
    
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            pad = 0
            if(filter_size[i] == 3):
                pad = 1
                
            conv2d = nn.Conv2d(in_channels, v, kernel_size=filter_size[i], padding=pad)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
        i += 1
    return nn.Sequential(*layers)

def net(network,in_channels=3):
    return make_layers(cfg[network],filter_size[network],batch_norm=True,in_channels=in_channels)`

Finally my code for training this model is as follows:

def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')

def load_checkpoint(filename='model_best.pth.tar'):
    return torch.load(filename)



def train_model(model, criterion, optimizer, scheduler,X,Y,num_epochs=25):

    since = time.time()

    best_model_wts = model.state_dict()
    best_acc = 0.0
    batch_size = 64
                                  
    for epoch in range(num_epochs):
        model.train(True)

        running_loss = 0.0
        running_corrects = 0

    for i in range(0,500000,batch_size):
        inputs,labels = X[i:i+batch_size], Y[i:i+batch_size]
        inputs,labels = Variable(inputs),Variable(labels)
        
       

        optimizer.zero_grad()
        outputs = model(inputs)
        _, preds = torch.max(outputs.data, 1)
        #y = Variable(torch.randn(28,1), requires_grad=False)
        
        loss = criterion(outputs,labels)
        #backward 
        loss.backward()
        optimizer.step()
        
        #statistics
        running_loss += loss.data[0]
        running_corrects += torch.sum(preds == labels.data)
        print('Loss is ', running_loss/(i+64))
        print('Accuracy is ' ,running_corrects/(i+64))


    dataset_sizes = 500000
    epoch_loss = running_loss / dataset_sizes
    epoch_acc = running_corrects / dataset_sizes
    print('{} Loss: {:.4f} Acc: {:.4f}'.format(epoch, epoch_loss, epoch_acc))

    # deep copy a model
    if epoch_acc > best_acc:
        
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'best_prec1': epoch_acc,
            'optimizer' : optimizer.state_dict(),
        },epoch_acc > best_acc)
        
        best_acc = epoch_acc
    
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:0f}s'.format(time_elapsed //60, time_elapsed%60))

print('Best val Acc: {:4f}'.format(best_acc))


load_checkpoint()

class_size = 5

X = np.load('./dataset_X.npy')
Y = np.load('./dataset_Y.npy')
Y = np.asarray([np.argmax(i) for i in Y])
X = torch.from_numpy(X).view(-1,1,28,28)
Y = torch.from_numpy(Y).type(torch.LongTensor)



model = SmallClassifier(utils.net('small_classifier',in_channels=1),class_size,batch_size=64)

criterion = nn.CrossEntropyLoss()
optimizer_ft = torch.optim.RMSprop(model.parameters())

# decay lr by facotr of 0.1 every 7 epochs

checkpoint = load_checkpoint()
model.load_state_dict(checkpoint['state_dict'])
optimizer_ft.load_state_dict(checkpoint['optimizer'])



exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
train_model(model,criterion,optimizer_ft,exp_lr_scheduler,X=X,Y=Y)

i have same problem
did you found el solution?