PyTorch Neural Networks

Create neural network class
class NeuralNet(nn.Module):
  def __init__(self):
    super().__init__()
    self.fc1 = nn.Linear(784, 128)
    self.fc2 = nn.Linear(128, 10)

  def forward(self, x):
    x = torch.relu(self.fc1(x))
    x = self.fc2(x)
    return x

model = NeuralNet().to(device)

Define loss function
criterion = nn.CrossEntropyLoss()
loss_mse = nn.MSELoss()

Define optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
sgd = torch.optim.SGD(model.parameters(), lr=0.01)

Training loop
for epoch in range(num_epochs):
  for batch_x, batch_y in train_loader:
    batch_x = batch_x.to(device)
    batch_y = batch_y.to(device)
    
    // Forward pass
    outputs = model(batch_x)
    loss = criterion(outputs, batch_y)
    
    // Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
  
  print(f"Epoch {epoch}, Loss: {loss.item()}")

Save model
torch.save(model.state_dict(), "model.pth")

Load model
model = NeuralNet()
model.load_state_dict(torch.load("model.pth"))
model.eval()

Save entire model
torch.save(model, "entire_model.pth")
model = torch.load("entire_model.pth")