yolov5 预训练模型 3分类-＞4分类

最新推荐文章于 2024-09-07 20:40:05 发布

highoooo

最新推荐文章于 2024-09-07 20:40:05 发布

阅读量1.1k

点赞数 1

分类专栏： Computer Vision

本文链接：https://blog.csdn.net/highoooo/article/details/120012330

版权

python

Computer Vision 专栏收录该内容

34 篇文章 3 订阅

订阅专栏

    # EMA
    ema = ModelEMA(model) if RANK in [-1, 0] else None
    # Resume
    start_epoch, best_fitness = 0, 0.0
    if pretrained:
        # Optimizer
        if ckpt['optimizer'] is not None:
            optimizer.load_state_dict(ckpt['optimizer'])
            best_fitness = ckpt['best_fitness']


        # EMA
        ckpt['new_ema'] = []
        for emaa in ckpt['ema'].state_dict():
            ckpt['new_ema'].append(emaa)
        new_weights = []
        for k,v in ckpt['ema'].float().state_dict().items():
            if k.startswith('model.24.m.0.weight'):
                new_v = torch.rand([27, 128, 1, 1])
                new_weights.append(new_v)
            elif k.startswith('model.24.m.1.weight'):
                new_v = torch.rand([27, 256, 1, 1])
                new_weights.append(new_v)
            elif k.startswith('model.24.m.2.weight'):
                new_v = torch.rand([27, 512, 1, 1])
                new_weights.append(new_v)
            elif k.startswith('model.24.m'):
                new_v = torch.rand([27])
                new_weights.append(new_v)
            else:
                new_weights.append(v)

        ckpt['my_weight'] = dict(zip(ckpt['new_ema'], new_weights))


        if ema and ckpt.get('ema'):
            ema.ema.load_state_dict(ckpt['my_weight'])
            # ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
            ema.updates = ckpt['updates']
            # ema.state_dict().update(ckpt['my_weight'])

默认加载方式

   # EMA
    ema = ModelEMA(model) if RANK in [-1, 0] else None

    # Resume
    start_epoch, best_fitness = 0, 0.0
    if pretrained:
        # Optimizer
        if ckpt['optimizer'] is not None:
            optimizer.load_state_dict(ckpt['optimizer'])
            best_fitness = ckpt['best_fitness']

        # EMA
        if ema and ckpt.get('ema'):
            ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
            ema.updates = ckpt['updates']