Pytorch与深度学习自查手册4-训练、可视化、日志输出、保存模型

训练和验证（包含可视化、日志、保存模型）

初始化模型、dataloader都完善以后，正式进入训练部分。

训练部分包括：

1. 及时的日志记录

2. tensorboard可视化log

3. 输入

4. 前向传播
5. loss计算
6. 反向传播

7. 权重更新

8. 固定步骤进行验证

9. 最佳模型的保存（+bad case输出）

日志记录

利用logging模块在控制台实时打印并及时记录运行日志。

from config import  *
import logging  # 引入logging模块
import os.path
class Logger:
    def __init__(self,mode='w'):
        # 第一步，创建一个logger
        self.logger = logging.getLogger()
        self.logger.setLevel(logging.INFO)  # Log等级总开关
        # 第二步，创建一个handler，用于写入日志文件
        rq = time.strftime('%Y%m%d%H%M', time.localtime(time.time()))
        log_path = os.getcwd() + '/Logs/'
        log_name = log_path + rq + '.log'
        logfile = log_name
        fh = logging.FileHandler(logfile, mode=mode)
        fh.setLevel(logging.DEBUG)  # 输出到file的log等级的开关
        # 第三步，定义handler的输出格式
        formatter = logging.Formatter("%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s")
        fh.setFormatter(formatter)
        # 第四步，将logger添加到handler里面
        self.logger.addHandler(fh)
        ch = logging.StreamHandler()
        ch.setLevel(logging.INFO)  # 输出到console的log等级的开关
        ch.setFormatter(formatter)
        self.logger.addHandler(ch)

完整的训练流程

import os
import math
import argparse
import torch
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
import torch.optim.lr_scheduler as lr_scheduler
import sys
from tqdm import tqdm
import torch

def train_one_epoch(model, optimizer, data_loader, device, epoch,tb_writer):
    model.train()
    loss_function = torch.nn.CrossEntropyLoss()
    mean_loss = torch.zeros(1).to(device)
    mean_acc = torch.zeros(1).to(device)
    optimizer.zero_grad()

    data_loader = tqdm(data_loader)
    for iteration, data in enumerate(data_loader):
        batch, labels = data
        pred = model(batch.to(device))

        loss = loss_function(pred, labels.to(device))
        loss.backward()
        mean_loss = (mean_loss * iteration + loss.detach()) / (step + 1)  # update mean losses
        pred = torch.max(pred, dim=1)[1]
        iter_acc=torch.eq(pred, labels.to(device)).sum()
		mean_acc+=iter_acc
        # 打印平均loss
        if iteration % 50 == 0:
        	data_loader.desc = "[epoch {}] mean loss {}".format(epoch, round(mean_loss.item(), 3))

        if not torch.isfinite(loss):
            print('WARNING: non-finite loss, ending training ', loss)
            sys.exit(1)

        optimizer.step()
        optimizer.zero_grad()
		tags=["train_loss","train_accuracy","learning_rate"]
        # tensorboard可视化
        for tag, value in zip(tags, [mean_loss.item(), iter_acc.item(), optimizer.param_groups[0]["lr"]]):
            tb_writer.add_scalars('Train\%s'%tag, value, iteration)
    return mean_loss.item(),mean_acc.item()


@torch.no_grad()
def evaluate(model, data_loader, device,best_acc=-1):
    model.eval()

    # 用于存储预测正确的样本个数
    sum_num = torch.zeros(1).to(device)
    # 统计验证集样本总数目
    num_samples = len(data_loader.dataset)

    # 打印验证进度
    data_loader = tqdm(data_loader, desc="validation...")
	bad_case=[]
    for step, data in enumerate(data_loader):
        batch, labels = data
        pred = model(batch.to(device))
        pred = torch.max(pred, dim=1)[1]
        tmp=torch.eq(pred, labels.to(device))
        sum_num += tmp.sum()
        bad_case.append((batch[~tmp],labels[~tmp]))
    # 计算预测正确的比例
    acc = sum_num.item() / num_samples
	if best_acc<acc:
        joblib.dump(bad_case,'bad_case.pkl')
    return acc

def main(args,logger):
    
    # 实例化模型
    model=Model()
    # 是否冻结权重
    if args.freeze_layers:
        print("freeze layers except fc layer.")
        for name, para in model.named_parameters():
            # 除最后的全连接层外，其他权重全部冻结
            if "fc" not in name:
                para.requires_grad_(False)

    pg = [p for p in model.parameters() if p.requires_grad]
    optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=0.005)
    
    lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf  # cosine
    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
    # 写入日志
    logger.logger.info('start training......\n')
    
	tb_writer = SummaryWriter(log_dir=args.log_dir)
	# 将模型写入tensorboard
	init_input = torch.zeros((1, 3, 224, 224), device=args.device)
	tb_writer.add_graph(model, init_input)
	
    best_acc=0#最佳模型的指标
	for epoch in range(args.epochs):
        mean_loss,mean_acc = train_one_epoch(model=model,
                                    optimizer=optimizer,
                                    data_loader=train_loader,
                                    device=device,
                                    epoch=epoch,
                                    tb_writer=tb_writer)
        # update learning rate
        scheduler.step()

        # validate
        acc = evaluate(model=model,
                       data_loader=val_loader,
                       device=device)

        # tensorboard可视化
        tb_writer.add_scalars('Validation\val_accuracy', acc, epoch)
        logger.logger.info('%d epoch train mean loss: %.2f \n'%(epoch,mean_loss))
        logger.logger.info('%d epoch train mean acc: %.2f \n'%(epoch,mean_acc))
        logger.logger.info('%d epoch validation acc: %.2f \n'%(epoch,acc))
        if epoch % args.save_epoch==0:
            checkpoint = {
    			'model_state_dict': model.state_dict(), #*模型参数
    			'optimizer_state_dict': optimizer.state_dict(), #*优化器参数
    			'scheduler_state_dict': scheduler.state_dict(), #*scheduler
    			'epoch': epoch,
    			'best_val_mae': best_valid_mae, 
    			'num_params': num_params
			}
			torch.save(checkpoint, os.path.join(args.save_dir, 'checkpoint-%d.pt'%epoch))
            logger.logger.info('save model %d successed......\n'epoch)
        # 保存最佳模型
        if best_acc<acc:
            best_acc=acc
            logger.logger.info('best model in %d epoch, train mean acc: %.2f \n'%(epoch,mean_acc))
        	logger.logger.info('best model in %d epoch, validation acc: %.2f \n'%(epoch,acc))
            checkpoint = {
    			'model_state_dict': model.state_dict(), #*模型参数
    			'optimizer_state_dict': optimizer.state_dict(), #*优化器参数
    			'scheduler_state_dict': scheduler.state_dict(), #*scheduler
    			'epoch': epoch,
    			'best_val_mae': best_valid_mae, 
    			'num_params': num_params
			}
			torch.save(checkpoint, os.path.join(args.save_dir, 'best_checkpoint.pt'))
            logger.logger.info('save best model  successed......\n')
        # 可视化图片预测结果
		# add figure into tensorboard
		fig = ...
		if fig is not None:
			tb_writer.add_figure("predictions vs. actuals",
								 figure=fig,
								 global_step=epoch)
        # 可视化权重不断更新的直方图
		# add conv1 weights into tensorboard
		tb_writer.add_histogram(tag="conv1",
								values=model.conv1.weight,
								global_step=epoch)
		tb_writer.add_histogram(tag="layer1/block0/conv1",
									values=model.layer1[0].conv1.weight,
									global_step=epoch)
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_classes', type=int, default=5)
    parser.add_argument('--epochs', type=int, default=30)
    parser.add_argument('--batch-size', type=int, default=16)
    parser.add_argument('--lr', type=float, default=0.001)
    parser.add_argument('--lrf', type=float, default=0.1)
    parser.add_argument('--save_epoch', type=float, default=3)
    parser.add_argument('--log_dir', type=float, default=3)
    # 数据集所在根目录
    data_root = "/home/data_set/"
    parser.add_argument('--data-path', type=str, default=img_root)

	#--freeze-layers #如果是True表示冻结除了全连接层以外的所有层的参数，在导入一些预训练的模型可以使用，可以加快模型训练
    parser.add_argument('--freeze-layers', type=bool, default=False)
    parser.add_argument('--device', default='cuda', help='device id (i.e. 0 or 0,1 or cpu)')

    opt = parser.parse_args()
	
	logger=Logger()
    main(opt,logger)

模型保存和断点继续训练

在训练模型过程中，对模型进行保存是很重要的。

核心包括两个内容：

1. 最优模型保存机制；
2. 如何从断点加载模型继续训练。

最优模型通常通过设置测试间隔，根据选定的指标，选择训练过程中表现最优的模型进行保存。需要保存以下内容到checkpoint：

1. 核心：模型参数、优化器参数、scheduler参数；
2. 其他：训练epoch、模型超参数、模型评价指标。

只保存和加载模型的话，还有其他方式可参考：PyTorch之保存加载模型 - 简书 (jianshu.com)

num_params = sum(p.numel() for p in model.parameters())
checkpoint = {
    'model_state_dict': model.state_dict(), #*模型参数
    'optimizer_state_dict': optimizer.state_dict(), #*优化器参数
    'scheduler_state_dict': scheduler.state_dict(), #*scheduler
    'epoch': epoch,
    'best_val_mae': best_valid_mae, 
    'num_params': num_params
}
torch.save(checkpoint, os.path.join(args.save_dir, 'checkpoint.pt'))

要实现断点继续训练只需要将模型上次保存的checkpoint加载进来，然后继续训练即可：

path_checkpoint = "./model_parameter/test/ckpt_best_50.pth"  # 断点路径
checkpoint = torch.load(path_checkpoint)  # 加载断点
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)
lr_schedule = ThelrscheduleClass(*args, **kwargs)
model.load_state_dict(checkpoint['model_state_dict'])  # 加载模型可学习参数

optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
start_epoch = checkpoint['epoch']  # 设置开始的epoch
lr_schedule.load_state_dict(checkpoint['lr_schedule'])#加载lr_scheduler

for epoch in range(start_epoch, args.epochs + 1):
    #……
    train_mae = train(model, device, train_loader, optimizer,scheduler, criterion_fn)
    #……

预测

@torch.no_grad()
def predict(model, data_loader, device):
    model.eval()

    # 用于存储预测正确的样本个数
    sum_num = torch.zeros(1).to(device)
    # 统计验证集样本总数目
    num_samples = len(data_loader.dataset)

    # 打印验证进度
    data_loader = tqdm(data_loader, desc="validation...")
	res=[]
    for step, batch in enumerate(data_loader):
        pred = model(batch.to(device))
        pred = torch.max(pred, dim=1)[1]
		res.extend(pred)
	return res