Write Epochs Results during Training to CSV File in PyTorch

When training machine learning models, tracking training progress across epochs - such as the loss value - is essential for debugging, visualization, and performance monitoring. One simple and effective way to do this is by saving results to a CSV file during training.

The provided PyTorch code trains a simple linear regression model and logs the training loss for each epoch to a CSV file. It uses the csv module to create a file named training.log and writes two columns: epoch and loss.

The model consists of a single linear layer trained with mean squared error (MSE) loss and stochastic gradient descent (SGD). During each of the 100 epochs, the code performs forward and backward passes, updates model parameters, and records the epoch number with the current loss value.

import torch
import csv

log_file = open('training.log', 'w', newline='')
writer = csv.writer(log_file)
writer.writerow(['epoch', 'loss'])

xs = torch.Tensor([[-2.0], [-1.0], [0.0], [1.0], [2.0], [3.0], [4.0]])
ys = torch.Tensor([[-3.0], [-1.0], [1.0], [3.0], [5.0], [7.0], [9.0]])

model = torch.nn.Sequential(
    torch.nn.Linear(1, 1)
)

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

epochs = 100
for epoch in range(epochs):
    optimizer.zero_grad()

    ys_pred = model(xs)

    loss = loss_fn(ys_pred, ys)
    loss.backward()
    optimizer.step()

    writer.writerow([epoch, loss.item()])

log_file.close()

After training, the CSV file contains a complete record of loss values over time, making it easy to analyze or visualize training progress later. Output example:

epoch,loss
0,35.87385177612305
1,28.88836097717285
2,23.28523826599121
3,18.79024314880371
4,15.183561325073242
..........
98,0.039155907928943634
99,0.03796852007508278