其他分享
首页 > 其他分享> > torch训练和推理的模板

torch训练和推理的模板

作者:互联网

def train(epoch):     model.train()     train_loss = 0     for data, label in train_loader:         data, label = data.cuda(), label.cuda()  # 将数据放入显卡         optimizer.zero_grad()         output = model(data)         loss = criterion(output, label)         loss.backward()         optimizer.step()         train_loss += loss.item()*data.size(0)    # loss需要取item()得到数字,才能相加     train_loss = train_loss/len(train_loader.dataset)     print('Epoch: {} \tTraining Loss: {:.6f}'.format(epoch, train_loss))       def val(epoch):           model.eval()     val_loss = 0     gt_labels = []     pred_labels = []     with torch.no_grad():         for data, label in test_loader:             data, label = data.cuda(), label.cuda()             output = model(data)             preds = torch.argmax(output, 1)             gt_labels.append(label.cpu().data.numpy())     # 将GPU中的tensor转化为numpy.array             pred_labels.append(preds.cpu().data.numpy())             loss = criterion(output, label)             val_loss += loss.item()*data.size(0)     val_loss = val_loss/len(test_loader.dataset)     gt_labels, pred_labels = np.concatenate(gt_labels), np.concatenate(pred_labels)     acc = np.sum(gt_labels==pred_labels)/len(pred_labels)     print('Epoch: {} \tValidation Loss: {:.6f}, Accuracy: {:6f}'.format(epoch, val_loss, acc))

标签:loss,torch,val,data,labels,label,train,推理,模板
来源: https://www.cnblogs.com/gagaein/p/16435856.html