train_epoch is a training routine that processes the training data once, batch by batch, computes loss, backpropagates gradients, and updates model parameters. Its exact work depends on the model: classification updates one network, while generative models may alternate or combine several objectives. section 2.4.2 section 7.4.2
A typical implementation follows this sequence:
- Read a batch from a data loader and move inputs and labels to the selected device. section 2.4.2
- Run the model, calculate its loss, clear accumulated gradients, backpropagate, and call the optimizer. section 2.4.2 section 7.4.2
- Accumulate a training-loss value and return or display an epoch-level result. section 2.4.2 section 7.4.2
def train_epoch():
tloss=0
for n,(imgs,labels) in enumerate(train_loader):
imgs=imgs.reshape(-1,28*28).to(device)
labels=labels.reshape(-1,).to(device)
preds=model(imgs)
loss=loss_fn(preds,labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
tloss+=loss.detach()
return tloss/n
The compact classifier implementation shows the core minibatch loop, tensor reshaping, loss calculation, gradient reset, backward pass, optimizer step, and loss accumulation. section 2.4.2
The routine is not a universal single-network template. CycleGAN trains two discriminators and two generators, detaches generated tensors for discriminator work, and uses automatic mixed precision. MuseGAN repeats critic updates before a generator update and adds a gradient penalty. A VAE combines reconstruction loss with KL divergence. section 6.4.1 section 13.5.1 section 7.4.2
Classification training loop
A multicategory image classifier calls train_epoch(), then val_epoch(), once per epoch. Training can run for up to 100 epochs, while validation loss controls early stopping after ten unimproved epochs. The trained model then assigns each image the class with the largest output using torch.argmax(). section 2.4.2
Listing 7.6 Defining the train_epoch() function
def train_epoch(epoch):
vae.train()
epoch_loss = 0.0
for imgs, _ in loader:
imgs = imgs.to(device)
mu, std, out = vae(imgs) #1
reconstruction_loss = ((imgs-out)**2).sum() #2
kl = ((std**2)/2 + (mu**2)/2 - torch.log(std) - 0.5).sum() #3
loss = reconstruction_loss + kl #4
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss+=loss.item()
print(f'at epoch {epoch}, loss is {epoch_loss}')
The VAE implementation makes the epoch boundary explicit while showing model mode selection, forward outputs, composite loss construction, and optimization. section 7.4.2
