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))