| import torch
|
|
|
|
|
| epochs = 0
|
| finetune = 2
|
|
|
|
|
| def train_model(model, data_loader, optimizer, criterion, device, epochs):
|
|
|
| model.train()
|
|
|
| for epoch in range(1, epochs + 1):
|
|
|
| total_loss = 0.0
|
| last_subepoch = -1
|
|
|
| for batch_idx, batch in enumerate(data_loader):
|
|
|
| x_batch, y_batch = batch
|
|
|
| x_batch = x_batch.to(device)
|
| y_batch = y_batch.to(device)
|
|
|
| optimizer.zero_grad()
|
|
|
| output = model(x_batch)
|
|
|
| output = output.view(-1, output.size(-1))
|
| y_batch = y_batch.view(-1)
|
|
|
| current_loss = criterion(output, y_batch)
|
| current_loss.backward()
|
|
|
| if current_loss.item() > 3:
|
| torch.nn.utils.clip_grad_norm_(
|
| model.parameters(),
|
| max_norm=0.5
|
| )
|
| else:
|
| torch.nn.utils.clip_grad_norm_(
|
| model.parameters(),
|
| max_norm=1.0
|
| )
|
|
|
| torch.cuda.empty_cache()
|
|
|
| optimizer.step()
|
|
|
| total_loss += current_loss.item()
|
|
|
|
|
| progress = (batch_idx + 1) / len(data_loader)
|
|
|
|
|
| subepoch = int(progress / 0.05) * 0.05
|
|
|
| if subepoch > last_subepoch and subepoch < 1.0:
|
| print(
|
| f"\rEpoch {epoch - 1 + subepoch:.2f}",
|
| end="",
|
| flush=True
|
| )
|
| last_subepoch = subepoch
|
|
|
| avg_loss = total_loss / len(data_loader)
|
|
|
|
|
| print("\r" + " " * 40 + "\r", end="")
|
|
|
|
|
| print(
|
| f"Epoch {epoch}/{epochs}, "
|
| f"Average Loss: {avg_loss:.4f}"
|
| )
|
|
|
| print("The weights have been updated.") |
