"""
对模型各个部分的一个调用工厂模块
"""
import torch

from .backbones import resnet_dcn
from .backbones import mnet
from .backbones import hrfeatnet
from .headers import common
from .framework import fairmot
from .framework import hrfairmot

_backbone_model_factory = {
    "resdcn": resnet_dcn.get_pose_net,
    "gridnet": mnet.get_backbone,
}

_header_model_factory = {
    "common": common.CommonHead,
}

def create_model(recipe, arch, heads_dict, head_conv):
    """
    创建整体模型
    Args:
        recipe: 模型类别
        arch: 模型的backbone架构
        heads_dict: 各个解耦头的配置
        head_conv: 解耦头中的卷积通道
    """
    print("==> Your recipe is:", recipe)
    arch_backbone = arch[:arch.find('_')] if '_' in arch else arch
    size = arch[arch.find('_') + 1:]
    get_model = _backbone_model_factory[arch_backbone]
    if arch_backbone not in ["gridnet"]:
        # resnet_dcn only, for simplicity
        num_layers = int(arch[arch.find('_') + 1:]) if '_' in arch else 0      
        feat_backbone = get_model(num_layers=num_layers)
    elif arch_backbone == "gridnet":
        if size == "tiny":
            feat_backbone = get_model(3, 32, 3, [1, 1, 1], 2)
        elif size == "normal":
            feat_backbone = get_model(3, 64, 3, [2, 2, 2], 2)
        elif size == "ball":
            feat_backbone = get_model(2, 32, 3, [1, 1, 1], 2, False)
        else:
            raise NotImplementedError
    else:
        raise NotImplementedError

    decouple_headers = _header_model_factory['common'](heads_dict, head_conv)

    if recipe == "fairmot":
        model = fairmot.FariMOT(feat_backbone, decouple_headers)
    elif recipe == "hrfairmot":
        # ??? 使用CNN进行增强训练速度更快有时候效果竟然还要更好 ???
        enhancer = hrfeatnet.FineFeatNet()   
        # print("Warning: code need to be refactored in [models/model_factory.py], hard coding is stupid")
        # enhancer = hrfeatnet.MAEFeatNet("/home/featurize/workspace/main/FSoccerMOT/models/mae_pretrain.pth")
        model = hrfairmot.HRFairMOT(feat_backbone, decouple_headers, enhancer)
    else:
        raise NotImplementedError
    return model

def load_model(
    model,
    model_path,
    optimizer=None,
    resume=False, 
    lr=None,
    lr_step=None,
    continuous = True,
):
    """
    根据已经训练好的权重加载模型, 而且在新创建一个模型并且还增加了新的模块的时候依然可以使用 只要根据shape对应上即可
    Args:
        model: 模型
        model_path: 模型保持的路径
        optimizer: 优化器
        resume: 可复用
        lr: 学习率
        lr_step: 衰减的步数
        continuous: 复用权重的同时继续上个训练任务 否则只是在复用权重
    """
    start_epoch = 0
    # 先将权重加载出来
    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
    print('loaded {}, epoch {}'.format(model_path, checkpoint['epoch']))
    state_dict_ = checkpoint['state_dict']
    state_dict = {}
  
    # convert data_parallal to model
    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]

    # 当前模型的初始化权重
    model_state_dict = model.state_dict()

    # check loaded parameters and created model parameters
    msg = 'If you see this, your model does not fully load the ' + \
            'pre-trained weight. Please make sure ' + \
            'you have correctly specified --arch xxx ' + \
            'or set the correct --num_classes for your own dataset.'

    # 遍历将权重加载到新的模型当中
    for k in state_dict:
        # 按照名称找到了对应的那个模块
        # 如果有同名的模块则判断shape是否相等然后加载进去
        if k in model_state_dict:                                       
            if state_dict[k].shape != model_state_dict[k].shape:
                print('Skip loading parameter {}, required shape{}, '\
                    'loaded shape{}. {}'.format(
                k, model_state_dict[k].shape, state_dict[k].shape, msg))
                # 如果shape不匹配则使用初始化权重
                state_dict[k] = model_state_dict[k]
        else:
            print('Drop parameter {}.'.format(k) + msg)
  
    # 反之整理那些没有存在于权重文件中的参数到新的state_dict当中
    for k in model_state_dict:
        if not (k in state_dict):
            print('No param {}.'.format(k) + msg)
            state_dict[k] = model_state_dict[k]

    # 使用非严格模式来加载已经训练的权重参数
    model.load_state_dict(state_dict, strict=False)

    # resume optimizer parameters
    # 这里同样加载了优化器的参数
    if optimizer is not None and resume and continuous:
        if 'optimizer' in checkpoint:
            optimizer.load_state_dict(checkpoint['optimizer'])
            start_epoch = checkpoint['epoch']
            start_lr = lr
        for step in 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 or start a new task.')
    if optimizer is not None:
        return model, optimizer, start_epoch
    else:
        return model

def save_model(
    path,
    epoch,
    model,
    optimizer=None
):
  """
  保存模型到文件 包括了轮数 模型以及优化器的参数等等
  """
  if isinstance(model, torch.nn.DataParallel):
    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)


"""
对模型各个部分的一个调用工厂模块
"""
import torch

from .backbones import resnet_dcn
from .backbones import mnet
from .backbones import hrfeatnet
from .headers import common
from .framework import fairmot
from .framework import hrfairmot

_backbone_model_factory = {
    "resdcn": resnet_dcn.get_pose_net,
    "gridnet": mnet.get_backbone,
}

_header_model_factory = {
    "common": common.CommonHead,
}

def create_model(recipe, arch, heads_dict, head_conv):
    """
    创建整体模型
    Args:
        recipe: 模型类别
        arch: 模型的backbone架构
        heads_dict: 各个解耦头的配置
        head_conv: 解耦头中的卷积通道
    """
    print("==> Your recipe is:", recipe)
    arch_backbone = arch[:arch.find('_')] if '_' in arch else arch
    size = arch[arch.find('_') + 1:]
    get_model = _backbone_model_factory[arch_backbone]
    if arch_backbone not in ["gridnet"]:
        # resnet_dcn only, for simplicity
        num_layers = int(arch[arch.find('_') + 1:]) if '_' in arch else 0      
        feat_backbone = get_model(num_layers=num_layers)
    elif arch_backbone == "gridnet":
        if size == "tiny":
            feat_backbone = get_model(3, 32, 3, [1, 1, 1], 2)
        elif size == "normal":
            feat_backbone = get_model(3, 64, 3, [2, 2, 2], 2)
        elif size == "ball":
            feat_backbone = get_model(2, 32, 3, [1, 1, 1], 2, False)
        else:
            raise NotImplementedError
    else:
        raise NotImplementedError

    decouple_headers = _header_model_factory['common'](heads_dict, head_conv)

    if recipe == "fairmot":
        model = fairmot.FariMOT(feat_backbone, decouple_headers)
    elif recipe == "hrfairmot":
        # ??? 使用CNN进行增强训练速度更快有时候效果竟然还要更好 ???
        enhancer = hrfeatnet.FineFeatNet()   
        # print("Warning: code need to be refactored in [models/model_factory.py], hard coding is stupid")
        # enhancer = hrfeatnet.MAEFeatNet("/home/featurize/workspace/main/FSoccerMOT/models/mae_pretrain.pth")
        model = hrfairmot.HRFairMOT(feat_backbone, decouple_headers, enhancer)
    else:
        raise NotImplementedError
    return model

def load_model(
    model,
    model_path,
    optimizer=None,
    resume=False, 
    lr=None,
    lr_step=None,
    continuous = True,
):
    """
    根据已经训练好的权重加载模型, 而且在新创建一个模型并且还增加了新的模块的时候依然可以使用 只要根据shape对应上即可
    Args:
        model: 模型
        model_path: 模型保持的路径
        optimizer: 优化器
        resume: 可复用
        lr: 学习率
        lr_step: 衰减的步数
        continuous: 复用权重的同时继续上个训练任务 否则只是在复用权重
    """
    start_epoch = 0
    # 先将权重加载出来
    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
    print('loaded {}, epoch {}'.format(model_path, checkpoint['epoch']))
    state_dict_ = checkpoint['state_dict']
    state_dict = {}
  
    # convert data_parallal to model
    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]

    # 当前模型的初始化权重
    model_state_dict = model.state_dict()

    # check loaded parameters and created model parameters
    msg = 'If you see this, your model does not fully load the ' + \
            'pre-trained weight. Please make sure ' + \
            'you have correctly specified --arch xxx ' + \
            'or set the correct --num_classes for your own dataset.'

    # 遍历将权重加载到新的模型当中
    for k in state_dict:
        # 按照名称找到了对应的那个模块
        # 如果有同名的模块则判断shape是否相等然后加载进去
        if k in model_state_dict:                                       
            if state_dict[k].shape != model_state_dict[k].shape:
                print('Skip loading parameter {}, required shape{}, '\
                    'loaded shape{}. {}'.format(
                k, model_state_dict[k].shape, state_dict[k].shape, msg))
                # 如果shape不匹配则使用初始化权重
                state_dict[k] = model_state_dict[k]
        else:
            print('Drop parameter {}.'.format(k) + msg)
  
    # 反之整理那些没有存在于权重文件中的参数到新的state_dict当中
    for k in model_state_dict:
        if not (k in state_dict):
            print('No param {}.'.format(k) + msg)
            state_dict[k] = model_state_dict[k]

    # 使用非严格模式来加载已经训练的权重参数
    model.load_state_dict(state_dict, strict=False)

    # resume optimizer parameters
    # 这里同样加载了优化器的参数
    if optimizer is not None and resume and continuous:
        if 'optimizer' in checkpoint:
            optimizer.load_state_dict(checkpoint['optimizer'])
            start_epoch = checkpoint['epoch']
            start_lr = lr
        for step in 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 or start a new task.')
    if optimizer is not None:
        return model, optimizer, start_epoch
    else:
        return model

def save_model(
    path,
    epoch,
    model,
    optimizer=None
):
  """
  保存模型到文件 包括了轮数 模型以及优化器的参数等等
  """
  if isinstance(model, torch.nn.DataParallel):
    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)


对模型各个部分的一个调用工厂模块 创建整体模型
     Args:
         recipe: 模型类别
         arch: 模型的backbone架构
         heads_dict: 各个解耦头的配置
         head_conv: 解耦头中的卷积通道 # resnet_dcn only, for simplicity # ??? 使用CNN进行增强训练速度更快有时候效果竟然还要更好 ??? # print("Warning: code need to be refactored in [models/model_factory.py], hard coding is stupid") # enhancer = hrfeatnet.MAEFeatNet("/home/featurize/workspace/main/FSoccerMOT/models/mae_pretrain.pth") 根据已经训练好的权重加载模型, 而且在新创建一个模型并且还增加了新的模块的时候依然可以使用 只要根据shape对应上即可
     Args:
         model: 模型
         model_path: 模型保持的路径
         optimizer: 优化器
         resume: 可复用
         lr: 学习率
         lr_step: 衰减的步数
         continuous: 复用权重的同时继续上个训练任务 否则只是在复用权重 # 先将权重加载出来 # convert data_parallal to model # 找到网络中真正存在权重的那一层而不是容器层 # 当前模型的初始化权重 # check loaded parameters and created model parameters # 遍历将权重加载到新的模型当中 # 按照名称找到了对应的那个模块 # 如果有同名的模块则判断shape是否相等然后加载进去 # 如果shape不匹配则使用初始化权重 # 反之整理那些没有存在于权重文件中的参数到新的state_dict当中 # 使用非严格模式来加载已经训练的权重参数 # resume optimizer parameters # 这里同样加载了优化器的参数 保存模型到文件 包括了轮数 模型以及优化器的参数等等