pytorch_training

training.py

# model object
model = MyModel()

# define lr, optimizer, loss function
lr = 0.001
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
loss_function = nn.CrossEntropyLoss()

# Train
epochs = 100
loss_his = []
for epoch in range(epochs):
    
    for data, labels in train_loader:  
        # Forward
        out = model(data)
        # Backward
        optimizer.zero_grad()
        loss = loss_function(out, labels)
        loss.backward()
        optimizer.step()
    
    loss_his.append(loss.item())