pytorch 0.4版本加载0.4.1 1.0更高版本的model

最新推荐文章于 2024-06-25 08:49:16 发布

芭蕉帘外雨声急

最新推荐文章于 2024-06-25 08:49:16 发布

阅读量2k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： CNN学习笔记

本文链接：https://blog.youkuaiyun.com/fly_egg/article/details/89469955

CNN学习笔记专栏收录该内容

2 篇文章

订阅专栏

本文介绍了一种从指定路径加载神经网络模型的方法，特别关注于如何处理DataParallel模型和过滤不必要的参数跟踪信息。通过使用PyTorch库，文章详细解释了如何自定义加载过程，确保模型与保存状态的一致性。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

def load_network(self, load_path, network, strict=True):
    if isinstance(network, nn.DataParallel):
        network = network.module

    model_dict = torch.load(load_path)
    filtered = {k: v for k, v in model_dict.items() if 'num_batches_tracked' not in k}
    network.load_state_dict(filtered, strict=strict)

    
    # network.load_state_dict(torch.load(load_path), strict=strict)