Computer Vision: model_1 poor results

[santiagorg2401]

Hello,
I am currently in the Computer Vision section and trained the model_1 using the GPU and functionalizing the train/test steps. However, I am getting poor results in training and testing accuracy. Here's the code.

class FashionMNISTModelV1(nn.Module):
  def __init__(self,
               input_shape: int,
               hidden_units: int,
               output_shape: int):
    super().__init__()
    self.layer_stack = nn.Sequential(
        nn.Flatten(),
        nn.Linear(input_shape, hidden_units),
        nn.ReLU(),
        nn.Linear(hidden_units, output_shape),
        nn.ReLU()
    )
  def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.layer_stack(x)

def train_step(model: torch.nn.Module,
               data_loader: torch.utils.data.DataLoader,
               loss_fn: torch.nn.Module,
               optimizer: torch.optim.Optimizer,
               accuracy_fn,
               device: torch.device):
  """Trains model with the dataset located on data_loader."""
  # Training loop
  train_loss, train_acc = 0, 0
  model.train()
  for batch, (X, y) in enumerate(data_loader):
    X, y = X.to(device), y.to(device)
    y_pred = model(X)

    loss = loss_fn(y_pred, y)
    train_loss += loss
    train_acc += accuracy_fn(y_pred.argmax(dim=1), y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

  # Calculate average loss and accuracy per epoch
  train_loss /= len(data_loader)
  train_acc /= len(data_loader)

  return train_loss, train_acc

def test_step(model: torch.nn.Module,
              data_loader: torch.utils.data.DataLoader,
              loss_fn: torch.nn.Module,
              accuracy_fn,
              device: torch.device):
  """Tests model with the dataset located on data_loader."""
  # Testing loop
  test_loss, test_acc = 0, 0
  model.eval()
  with torch.inference_mode():
    for X, y in data_loader:
      X, y = X.to(device), y.to(device)
      test_pred = model(X)

      test_loss += loss_fn(test_pred, y)
      test_acc += accuracy_fn(test_pred.argmax(dim=1), y)

    test_loss /= len(data_loader)
    test_acc /= len(data_loader)

  return test_loss, test_acc

EPOCHS = 3; LOG_STEP = 1; LEARNING_RATE = 42

# Initialize model
torch.manual_seed(SEED); torch.cuda.manual_seed(SEED)
model_1 = FashionMNISTModelV1(784,
                              10,
                              len(class_names)).to(device)

# Create a loss function, an optimizer and a metrics function
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model_1.parameters(),
                            lr=LEARNING_RATE)
accuracy_fn = Accuracy(task="multiclass",
                       num_classes=len(class_names)).to(device)

model_1_t0 = timer()

# Use tqdm progress bar
for epoch in tqdm(range(EPOCHS)):
  print(f"-------------Epoch: {epoch}-------------")
  # Training loop
  train_loss, train_acc = train_step(model=model_1,
                                     data_loader=train_dataloader,
                                     loss_fn=loss_fn,
                                     optimizer=optimizer,
                                     accuracy_fn=accuracy_fn,
                                     device=device)

  # Testing loop
  test_loss, test_acc = test_step(model=model_1,
                                  data_loader=test_dataloader,
                                  loss_fn=loss_fn,
                                  accuracy_fn=accuracy_fn,
                                  device=device)

  # Log experiment data
  if epoch % LOG_STEP == 0 or epoch == EPOCHS - 1:
    m = f"Epoch: {epoch} | " \
        f"Loss: {train_loss:.5f} | Accuracy: {100*train_acc:.2f}%" \
        f"| Test loss: {test_loss:.5f} | Test accuracy: {100*test_acc:.2f}%"
    print(m)

model_1_t1 = timer()
model_1_time = print_train_time(start=model_1_t0,
                                end=model_1_t1,
                                device=device)

It prompts:

-------------Epoch: 0-------------
Epoch: 0 | Loss: 137.63127 | Accuracy: 10.02%| Test loss: 2.30259 | Test accuracy: 9.99%
-------------Epoch: 1-------------
Epoch: 1 | Loss: 2.30256 | Accuracy: 10.00%| Test loss: 2.30259 | Test accuracy: 10.01%
-------------Epoch: 2-------------
Epoch: 2 | Loss: 2.30256 | Accuracy: 10.00%| Test loss: 2.30259 | Test accuracy: 9.99%
Train time on cuda: 40.476 seconds.

And using the eval function:

{'model_name': 'FashionMNISTModelV1',
 'model_loss': 2.30259108543396,
 'model_acc': 0.00099940097425133}

[sametucuncu]

LEARNING_RATE should be 0.1

[santiagorg2401]

Did I have a good laugh after realizing. Thanks for the observation!