| import torch.optim as optim |
| import torch.nn as nn |
|
|
| class Trainer: |
| def __init__(self, model): |
| self.model = model |
| self.criterion = nn.CrossEntropyLoss() |
| self.optimizer = optim.SGD(self.model.parameters(), lr=0.01) |
|
|
| def train(self, train_loader, epochs=5): |
| for epoch in range(epochs): |
| for images, labels in train_loader: |
| self.optimizer.zero_grad() |
| outputs = self.model(images.view(-1, 28*28)) |
| loss = self.criterion(outputs, labels) |
| loss.backward() |
| self.optimizer.step() |
|
|
| print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}") |
