Pytorch踩坑记录之模型载入

最新推荐文章于 2022-03-20 00:35:33 发布

ReLuJie

最新推荐文章于 2022-03-20 00:35:33 发布

阅读量875

点赞数 1

分类专栏：深度学习文章标签： pytorch load_state_dict

本文链接：https://blog.csdn.net/On_theway10/article/details/90518743

版权

深度学习专栏收录该内容

66 篇文章 1 订阅

订阅专栏

def load_state_dict(self, state_dict):
        ''' Self-write load state_dict '''

        for name, param in state_dict.items():
            if name in self._channel_dict:
                if 'bn' in name:
                    param = param.unsqueeze(0).data
                else:
                    param = param.data
                try:
                    self._channel_dict[name].copy_(param)
                except Exception:
                    raise RuntimeError('While copying the parameter named {}, '
                                       'whose dimensions in the model are {} and '
                                       'whose dimensions in the checkpoint are {}.'
                                       .format(name, self._channel_dict[name].size(), param.size()))

refer-link : https://blog.csdn.net/a137376864/article/details/78654618