train loss代码

以下是一个简单的PyTorch代码段，用于计算训练集上的平均损失（loss）。

import torch.nn.functional as F

def train(model, optimizer, criterion, train_loader):
    model.train() # 将模型设置为训练模式
    train_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad() # 清除梯度
        output = model(data) # 前向传播
        loss = criterion(output, target) # 计算损失
        loss.backward() # 反向传播
        optimizer.step() # 更新权重
        train_loss += loss.item() # 累加损失
    train_loss /= len(train_loader.dataset) # 计算平均损失
    return train_loss

在这个代码段中，我们首先将模型设置为训练模式，然后遍历训练数据集中的每个batch。在每个batch中，我们执行以下操作：

1. 清除梯度
2. 前向传播
3. 计算损失
4. 反向传播
5. 更新权重
6. 累加损失

最后，我们计算整个训练集上的平均损失。这个平均损失可以用来监控模型的训练过程，并且可以用于比较不同模型的表现。