[深度学习]-训练和测试模型的常用代码

数据转换

data2 = torch.tensor(data1)

非训练数据（训练结束后的tensor）的转换：
data2 = np.array(data1)
如果是训练过程中需要转换，则：
data2 = data1.cpu().data.numpy() 动态，会一起改变
data2 = data1.cpu().detach().numpy() 静态

如果第一维，维度为1，可以去掉：（1，8，8）-> （8，8）
data2 = np.concatenate(datat1, axis=0)
如果包含维度为1的维度，可以去掉：（1，8，8，1）-> （8，8）
data2 = data1.squeeze()
转换维度层数：（1，8，3）-> （8，1，3）
data2 = np.swapaxes(1, 0, 2)

加载数据集dataloader

from torch.utils.data import DataLoader
form 自己写的dataset import Dataset

train_set = Dataset(train=True)
val_set = Dataset(train=False)

image_datasets = {
    'train': train_set, 'val': val_set
}

batch_size = 4

dataloaders = {
    'train': DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2),
    'val': DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=2)
}

dataset_sizes = {
    x: len(image_datasets[x]) for x in image_datasets.keys()
}
print(dataset_sizes)

for epoch in range(num_epochs):
    for phase in ['train', 'val']:
        if phase == 'train':
            # for param_group in optimizer.param_groups:
                # print("LR", param_group['lr'])
            model.train()
        else:
            model.eval()

以上适用于train一遍test一遍的情况

或者分别加载训练和测试：

train_dataset = Dataset('train')
train_data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True,
                                                    num_workers=2, collate_fn=collate_fn)

test_dataset = Dataset('eval')
test_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8, shuffle=False,
                                                   num_workers=2, collate_fn=collate_fn)

自己写Dataset

from torch.utils.data import Dataset
import os
import cv2
import torch
import numpy as np

class Dataset(Dataset):
    def __init__(self,train):
        if train:
            self.datapath = {'image': '/home/myy/code/Final_Project/data_train.txt', 'target':'/home/myy/code/Final_Project/gt_train.txt'}
        else:
            self.datapath = {'image': '/home/myy/code/Final_Project/data_test.txt', 'target':'/home/myy/code/Final_Project/gt_test.txt'}
            # self.datapath = {'image': '/home/myy/code/Final_Project/test_small_data.txt', 'target':'/home/myy/code/Final_Project/test_small.txt'}
        self.image_list, self.target_list = self.read_txt(self.datapath)

此处可以依据需要自己定义一些函数
注意调用前要加上self.
比如以下两个读取数据的函数，read_txt、read_json就是自己定义的
    def read_txt(self,datapath):
        im =[]
        target_image = []
        print(datapath)
        with open(datapath['image'], 'r') as f:
            image_list = f.readlines()
        with open(datapath['target'], 'r') as f:
            target_list = f.readlines()
        return image_list, target_list

    def read_json(save_path, encoding='utf8'):
        jsondata = []
        with open(save_path, 'r', encoding=encoding) as f:
            content = f.read()
            content = json.loads(content)
            for key in content:
                jsondata.append(content[key])
            return jsondata

    def __getitem__(self, item):
        # 最核心的部分，经过处理，要返回输入和gt

        return img, target

    def __len__(self):
        # 这可以根据具体情况修改，不写也行
        return len(self.data)

模型训练

torch.save(model.state_dict(), model_path)

import copy
best_model_wts = copy.deepcopy(model.state_dict())

trainable_vars = list(model.roi_heads.box_predictor.parameters())+ \
                 list(model.roi_heads.box_head.parameters())+ \
                 list(model.rpn.head.parameters())
                #  list(model.backbone.body.parameters())+ \
                #  list(model.backbone.fpn.parameters())+ \
 optimizer = torch.optim.SGD(trainable_vars , lr=1e-3, momentum=0.5)

模型测试

主要代码

device = torch.device("cuda:0" if  torch.cuda.is_available() else "cpu")
print(device)
num_class = 1
model_path = '训练模型保存路径/保存名称.pth'
model = 模型(num_class).to(device)
model.load_state_dict(torch.load(model_path))
model.eval()

import cv2
from torchvision import transforms, datasets, models
from torch.utils.data import  DataLoader
import torch
import numpy as np
import os
from sklearn import metrics
import matplotlib.pyplot as plt

device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
print(device)
num_class = 3
model_path = 模型路径
model = 模型(num_class).to(device)
model.load_state_dict(torch.load(model_path))
model.eval()   # Set model to evaluate mode

test_dataset = 数据集读取(train=False)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=2)

def turn(l):
    l = l.data.cpu().numpy()
    l = l.squeeze()
    l = np.swapaxes(l, 0, 2)
    l = np.swapaxes(l, 0, 1)
    return l

for inputs, labels in test_loader:
    model.to(device)
    inputs = inputs.to(device)
    labels = labels.to(device)

    pred = model(inputs)
    # pred = torch.relu(pred)
    pred = turn(pred)
    gt = turn(labels)

评价指标

以分割为例，经过.flatten()处理。

def acc(pred, gt):
    tp = 0
    tn = 0
    fp = 0
    fn = 0
    num = len(pred)
    for i in range(num):
        if pred[i] > 0 and gt[i] == 1:
            tp += 1
        if pred[i] > 0 and gt[i] == 0:
            fp += 1
        if pred[i] == 0 and gt[i] == 1:
            fn += 1
        if pred[i] == 0 and gt[i] == 0:
            tn += 1
    acc = (tp + tn) / num
    iou = tp / (tp + fp + fn)
    rec = tp / (tp + fn)
    pre = tp / (tp + fp)
    f1 = 2 * pre * rec / (pre + rec)
    print("mAcc is :{},  mIou is :{}, recall is :{}, precision is :{}, f1 is :{}".format(acc, iou, rec, pre, f1))

def draw_roc(pred, gt, name):
    tpr, fpr, thresholds = metrics.roc_curve(gt, pred, pos_label=0)
    plt.figure
    plt.plot(fpr, tpr, label = name)
    plt.xlabel('FPR')
    plt.ylabel('TPR')
    plt.legend(loc = 'lower right')
    plt.title(name)
    plt.savefig('路径/{}.png'.format(name))
    # plt.close()  如果有多个类别，不close()就会画在一张图上

Original: https://www.cnblogs.com/camilia/p/16000412.html
Author: CAMILIA
Title: [深度学习]-训练和测试模型的常用代码