Unable to train CNN model on CIFAR-10

average_ml_enjoyer: <p>Hi guys! I wrote a bit of code on training the CNN model on CIFAR-10, but the model seems to not learn, jumping around an accuracy of 10%. I borrowed the hyperparameters from another CNN model on the same dataset that seem to work, but for some reason my model just does not train itself. Can anyone tell me what the problem is?<br>
*this is my first time asking on forums, so if there is any problem with uploading code or what should be posted, please inform me((</p>
<pre><code class="lang-auto">import pandas as pd

import torch
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import torch.nn as nn
import torch.optim as optim
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import TensorDataset, DataLoader, Subset


# Check if CUDA is available
cuda_available = torch.cuda.is_available()
print(f"CUDA available: {cuda_available}")

# Print the name of the GPU
if cuda_available:
    device="cuda"
    print(f"GPU: {torch.cuda.get_device_name(0)}")
else:
    device="cpu"
#import data
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
temp_data = torchvision.datasets.CIFAR10(root="data",train=True,download=True,transform=transform)
test_data = torchvision.datasets.CIFAR10(root="data",train=False,download=True,transform=transform)
train_indices, val_indices = train_test_split(list(range(len(temp_data))), test_size=0.2, shuffle=True)
train_data = Subset(temp_data, train_indices)
val_data = Subset(temp_data, val_indices)
train_dataloader = DataLoader(train_data, batch_size=4, shuffle = True)
val_dataloader = DataLoader(val_data, batch_size=len(val_data), shuffle=True)
test_dataloader = DataLoader(val_data, batch_size=len(test_data), shuffle=True)
# Get a batch of data from the DataLoader
for batch in val_dataloader:
    x_val, y_val = batch
x_val, y_val = x_val.to(device), y_val.to(device)
#CNN
class model(nn.Module):
  def __init__(self):
    super(model, self).__init__()
    self.conv1 = nn.Sequential(
      nn.Conv2d(
          in_channels = 3,
          out_channels = 16,
          kernel_size = 3,
          stride = 1,
          padding = 1,
      ),
      nn.BatchNorm2d(16),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size = 2)
    )
    self.conv2 = nn.Sequential(
      nn.Conv2d(
          in_channels = 16,
          out_channels = 32,
          kernel_size = 3,
          stride = 1,
          padding = 1,
      ),
      nn.BatchNorm2d(32),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size = 2)

    )
    self.out = nn.Sequential(
        nn.Linear(32*8*8, 256),
        nn.ReLU(),
        nn.Linear(256, 10)
        )
  def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)
    #print(x.shape)
    x = x.view(-1, 32*8*8)
    x = self.out(x)
    return x
model = model()
model = model.to(device)
#hyperparameters
learning_rate = 0.001
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
epochs = 5
train_loss = []
vali_loss = []
test_loss_list = []
loss_fn = nn.CrossEntropyLoss()
#train loop
def train_loop(dataloader, model, loss_fn, optimizer):
  size = len(dataloader.dataset)
  model.train()
  for batch, (x, Y) in enumerate(train_dataloader):
    x, Y = x.to(device), Y.to(device)
    pred = model(x)

    loss = loss_fn(pred, Y)
    loss.backward()
    optimizer.zero_grad()
    optimizer.step()

    if batch%100 == 0:
      model.eval()
      val_loss = loss_fn(model(x_val), y_val).item()
      current = (batch + 1) * len(x)
      print(f"loss:{loss:&gt;7f}[{current:&gt;5d}/{size:&gt;5d}]")
      train_loss.append(loss.item())
      vali_loss.append(val_loss)

#test loop
def test_loop(dataloader, model, loss_fn, optimizer):
  model.eval()
  count = 0
  test_loss = 0
  size = 0
  num_batches = len(dataloader)

  with torch.no_grad():
   for x, y in dataloader:
    pred = model(x)
    x, y = x.to(device), y.to(device)

    test_loss += loss_fn(pred, y).item()
    count += (pred.argmax(1) == y).sum().item()
    size += y.size(0)
  test_loss = test_loss / num_batches
  accuracy = count / size
  print(f"Test Loss: {test_loss:.4f}, Accuracy: {accuracy:.4f}")
  test_loss_list.append(test_loss)

for x in range(epochs):
  print(f"epoch{x+1}\n")
  train_loop(train_dataloader, model, loss_fn, optimizer)
  test_loop(test_dataloader, model, loss_fn, optimizer)
plt.plot(train_loss, label = "train loss")
plt.plot(test_loss_list, label = "test loss")
plt.plot(vali_loss, label = "vali loss")
plt.xlabel("epoch")
plt.ylabel("loss")
plt.title("learning curve")
plt.legend()
plt.show()

</code></pre>

thebrownfrog: <p><code>optimizer.zero_grad()</code> is ran just before <code>optimizer.step()</code>, which makes the optimizer unable to work. Run <code>optimizer.zero_grad()</code> before computing loss (before line <code>loss = loss_fn(pred, Y)</code>). A chatbot explained it nicely: “The purpose of zero_grad() is to clear the gradients accumulated during the backward pass so that they don’t accumulate across iterations”</p>
<p>*first time answering</p>

average_ml_enjoyer: <p>wow it worked! thanks&lt;3</p>