【CenterFusion】模型的创建、导入、保存CenterFusion/src/lib/model/model.py

时间:2024-03-21 19:34:07
from __future__ import absolute_import from __future__ import division from __future__ import print_function import torchvision.models as models import torch import torch.nn as nn import os from .networks.dla import DLASeg from .networks.resdcn import PoseResDCN from .networks.resnet import PoseResNet from .networks.dlav0 import DLASegv0 from .networks.generic_network import GenericNetwork _network_factory = { 'resdcn': PoseResDCN, 'dla': DLASeg, 'res': PoseResNet, 'dlav0': DLASegv0, 'generic': GenericNetwork } def create_model(arch, head, head_conv, opt=None): num_layers = int(arch[arch.find('_') + 1:]) if '_' in arch else 0 ''' 处理字符串 arch = dla_34 ,将下划线后半部分取出 最后 num_layers = 34 ''' arch = arch[:arch.find('_')] if '_' in arch else arch ''' 将 arch = dla_34 中下划线前半部分取出 最后 arch = 'dla' ''' model_class = _network_factory[arch] ''' 根据 arch = 'dla' 获取 _network_factory 中的值 最后 model_class = DLASeg DLASeg 类定义在 CenterFusion/src/lib/model/networks/dla.py 第 594 行 ''' model = model_class(num_layers, heads=head, head_convs=head_conv, opt=opt) ''' 配置模型 ''' return model def load_model(model, model_path, opt, optimizer=None): start_epoch = 0 ''' 设定初始轮次 = 0 ''' checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) print('loaded {}, epoch {}'.format(model_path, checkpoint['epoch'])) ''' torch.load() 函数:用来加载 torch.save() 保存的模型文件 ''' state_dict_ = checkpoint['state_dict'] ''' 获取 checkpoint 模型文件中的 state_dict 属性 这个属性存放训练过程中需要学习的权重和偏执系数 state_dict 作为 python 的字典对象将每一层的参数映射成 tensor 张量 需要注意的是 torch.nn.Module 模块中的 state_dict 只包含卷积层和全连接层的参数 ''' state_dict = {} for k in state_dict_: if k.startswith('module') and not k.startswith('module_list'): state_dict[k[7:]] = state_dict_[k] else: state_dict[k] = state_dict_[k] ''' startswith(str) 函数:检测字符串 str,检测到返回 True,否则返回 False 这里只执行了 else 语句,相当于保存导入模型的网络参数 ''' model_state_dict = model.state_dict() ''' 浅拷贝 main.py 中创建的新模型 DLA 的网络参数 ''' for k in state_dict: ''' 遍历导入的模型中的每层网络参数 ''' if k in model_state_dict: ''' 判断新模型的网络参数中是否有导入的模型的参数 是有的,因为导入的模型也是 DLA 模型 ''' if (state_dict[k].shape != model_state_dict[k].shape) or \ (opt.reset_hm and k.startswith('hm') and (state_dict[k].shape[0] in [80, 1])): ''' 第一个条件为 True 其余条件全部为 False ''' if opt.reuse_hm: ''' 不执行 ''' print('Reusing parameter {}, required shape{}, '\ 'loaded shape{}.'.format( k, model_state_dict[k].shape, state_dict[k].shape)) # todo: bug in next line: both sides of < are the same if state_dict[k].shape[0] < state_dict[k].shape[0]: model_state_dict[k][:state_dict[k].shape[0]] = state_dict[k] else: model_state_dict[k] = state_dict[k][:model_state_dict[k].shape[0]] state_dict[k] = model_state_dict[k] elif opt.warm_start_weights: ''' 不执行 ''' try: print('Partially loading parameter {}, required shape{}, '\ 'loaded shape{}.'.format( k, model_state_dict[k].shape, state_dict[k].shape)) if state_dict[k].shape[1] < model_state_dict[k].shape[1]: model_state_dict[k][:,:state_dict[k].shape[1]] = state_dict[k] else: model_state_dict[k] = state_dict[k][:,:model_state_dict[k].shape[1]] state_dict[k] = model_state_dict[k] except: print('Skip loading parameter {}, required shape{}, '\ 'loaded shape{}.'.format( k, model_state_dict[k].shape, state_dict[k].shape)) state_dict[k] = model_state_dict[k] else: ''' 执行该 else 中的语句 ''' print('Skip loading parameter {}, required shape{}, '\ 'loaded shape{}.'.format( k, model_state_dict[k].shape, state_dict[k].shape)) state_dict[k] = model_state_dict[k] ''' 将新模型的网络参数赋值给导入的模型中 ''' else: print('Drop parameter {}.'.format(k)) for k in model_state_dict: if not (k in state_dict): print('No param {}.'.format(k)) state_dict[k] = model_state_dict[k] ''' 给导入的模型添加没有的参数 ''' model.load_state_dict(state_dict, strict=False) ''' 使用 state_dict 反序列化模型参数字字典,用来加载模型参数 将 state_dict 中的 parameters 和 buffers 复制到此 module 及其子节点中 简述:给模型对象加载训练好的模型参数,即加载模型参数 ''' #冻结骨干网,没有执行 if opt.freeze_backbone: for (name, module) in model.named_children(): if name in opt.layers_to_freeze: for (name, layer) in module.named_children(): for param in layer.parameters(): param.requires_grad = False # 恢复优化器参数,没有执行 if optimizer is not None and opt.resume: if 'optimizer' in checkpoint: start_epoch = checkpoint['epoch'] start_lr = opt.lr for step in opt.lr_step: if start_epoch >= step: start_lr *= 0.1 for param_group in optimizer.param_groups: param_group['lr'] = start_lr print('Resumed optimizer with start lr', start_lr) else: print('No optimizer parameters in checkpoint.') if optimizer is not None: ''' 执行该 if 语句 ''' return model, optimizer, start_epoch else: return model def save_model(path, epoch, model, optimizer=None): if isinstance(model, torch.nn.DataParallel): ''' isinstance(object, classinfo) 判断一个函数 object 是否是一个已知的类型 classinfo 是则返回 True,反之返回 False ''' state_dict = model.module.state_dict() else: state_dict = model.state_dict() ''' 获取模型的参数矩阵 ''' data = {'epoch': epoch, 'state_dict': state_dict} if not (optimizer is None): data['optimizer'] = optimizer.state_dict() ''' 获取模型的优化器 ''' torch.save(data, path) ''' 保存模型 '''