def train(model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = nn.functional.nll_loss(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()))
时间: 2024-03-14 22:49:03 浏览: 211
example.train
这段代码实现了一个简单的训练函数,输入参数包括模型、设备、训练数据加载器、优化器和训练轮数。函数中通过迭代训练数据加载器中的每一个batch,将数据和标签移动到设备上,然后将优化器的梯度清零,运行模型并计算损失,然后进行反向传播和梯度更新。每100个batch打印一次训练进度。
需要注意的是,损失函数使用的是负对数似然损失函数(Negative Log-Likelihood Loss),即nn.functional.nll_loss。这个损失函数通常用于多分类问题,需要将模型的输出先通过softmax函数转换为概率分布,然后计算对数似然损失。
此外,模型训练的过程还需要考虑其他因素,比如学习率的调整、模型的保存等等。
阅读全文