PyTorch实用模块总结

1. 使用预训练的模型初始化网络的一部分参数

# 1. 本地加载预训练模型
pretrained_dict = torch.load(weights_path)
# 或:url加载方式
pretrained_dict = torch.utils.model_zoo.load_url(weights_url)

# 2. 获取网络的state_dict
model_dict = model.state_dict()

# 3. 创建新的state_dict
new_state_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 4. 加载新的state_dict
model.load_state_dict(new_state_dict)

2. 调用.cuda()非常慢

解决方案:采用source安装或者pip安装,不要用conda安装
参考自:
[1] https://github.com/pytorch/pytorch/issues/537
[2] https://discuss.pytorch.org/t/model-cuda-takes-long-time/102/11

你可能感兴趣的:(PyTorch)