MNIST code problem

I modify and combine the code on the internet recently ,
Did it seem to be wrong about train auccuracy , train loss, test auccuracy , test auccuracy?
or others improve suggestion?

thanks

code:

 import torch
 import torch.nn as nn
 import matplotlib.pyplot as plt
 from torch.utils.data import DataLoader
 from torchvision import datasets, transforms

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

train_file = datasets.MNIST(root='./data/',train=True,transform=transforms.ToTensor(),download=True)

test_file = datasets.MNIST(root='./dataset/',train=False,transform=transforms.ToTensor(),download=True)

train_loader = DataLoader(dataset=train_file,batch_size=128,shuffle=True)

test_loader = DataLoader(dataset=test_file,batch_size=128,shuffle=False)
class CNN(nn.Module):

    def __init__(self):

        super(CNN,self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Linear(64 * 7 * 7, 10)

    def forward(self,x):
        x = self.conv(x)
        x = x.view(x.size(0), -1) 
        y = self.fc(x)
        return y

model = CNN().to(device)
print(model)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3,momentum=0.9)
loss_fn = nn.CrossEntropyLoss()
def train(dataloader , model ,loss_fn, optimizer):

    size = len(dataloader.dataset)

    num_batches = len(dataloader)

    train_loss = 0

    train_Acc = 0

    for data , targets in train_loader:

        data = data.to(device)

        targets = targets.to(device)

        # Compute prediciton error

        pred = model(data)

        loss = loss_fn(pred,targets)

        # Backpropagation

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

   

        train_loss += loss.item()

        train_Acc += (pred.argmax(1) == targets).type(torch.float).sum().item()

    train_loss /= num_batches

    train_Acc /= size

    print('Train Auccuracy: %.1f , Train loss: %.4f \n' %(100*train_Acc , train_loss) )
def test(dataloader, model ,loss_fn):

    size = len(dataloader.dataset)

    num_batches = len(dataloader)

    model.eval()

    test_loss = 0 

    test_acc = 0

    with torch.no_grad():

        for data, targets in test_loader:

            data,targets = data.to(device),targets.to(device)

            pred = model(data)

            test_loss += loss_fn(pred,targets).item()

            test_acc += (pred.argmax(1) == targets).type(torch.float).sum().item()

    test_loss /= num_batches

    test_acc /= size

    print('Test Auccuracy: %.1f , Test loss: %.4f \n' %(100*test_acc , test_loss) )
epochs = 5

for t in range(epochs):

    print(f"Epoch {t+1}\n-------------------------------")

    train(train_loader, model, loss_fn, optimizer)

    test(test_loader, model, loss_fn)

print("Done!")

Can’t really help without actually looking at the error or outputs.
However in your train, test function you are iterating over train_loader, test_loader and not the parameter dataloader

@arya47

outputs:

Epoch 1
-------------------------------
Train Auccuracy: 73.8 , Train loss: 1.0836 

Test Auccuracy: 90.2 , Test loss: 0.3438 

Epoch 2
-------------------------------
Train Auccuracy: 91.2 , Train loss: 0.3006 

Test Auccuracy: 92.7 , Test loss: 0.2449 

Epoch 3
-------------------------------
Train Auccuracy: 93.5 , Train loss: 0.2227 

Test Auccuracy: 94.7 , Test loss: 0.1786 

Epoch 4
-------------------------------
Train Auccuracy: 94.8 , Train loss: 0.1761 

Test Auccuracy: 95.8 , Test loss: 0.1449 

Epoch 5
-------------------------------
Train Auccuracy: 95.7 , Train loss: 0.1468
Test Auccuracy: 96.3 , Test loss: 0.1235 

Done!

Could you tell me how to do parameter dataloader?

By parameter dataloader i mean note in your test(dataloader, model, loss_fn) you have not used dataloader in the function instead used test_loader. This should not cause an issue but just wanted to point that out.

Looking at your outputs it seems to be working. In the future to check if your model is training you can try to overfit on a small dataset of say 50 samples.

oh, you mean
in the train_loader , I use the optimizer to parameter dataloader.
in the test_loader , I don’t use any method to do parameter dataloader.
So it seems to be problematic?

but the ouputs can work normally?

Nah i was referring to the train function
def train(dataloader, model, loss_fn):

In the function you have used train_loader (for data, target in train_loader) it should be (for data, target in train_loader).
This wont make a difference if you run on the same kernel

oh. I got it .
I should use it :
with torch.no_grad():
the same method on the train function?

torch.no_grad() should only be used in testing.

yes , so this code is feasible?