PyTorch的module的state_dict函数介绍

在PyTorch中，torch.nn.Module模块中的state_dict函数获得的一个字典变量，其存放训练过程中需要学习的权重和偏执系数，如下所示：

代码1：

import torch.nn as nn

module = nn.Linear(2, 2)

print(module.state_dict().keys())

代码2：

import torch.nn as nn
import torch.nn.functional as F

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))


module = MyModel()

print(module.state_dict().keys())

参考

https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.state_dict

学习建议

我在工作中接触了很多人，有学生有上班族，尤其是学生，他们的动手能力相当的弱，不是他们不努力，而是方向跑偏了，很多人执着于读论文、跑开源、学算法，但是忽略了最底层的基本功，连PyTorch都用不好，这是一个非常错误的做法。很多人问我，某某你会吗？我总数回答：只要是用PyTorch，我都会。因为我深刻体会到了学好PyTorch，我也提醒各位读者，学好PyTorch十分重要。欢迎大家关注红黑树技术联盟，这里有高质量的PyTorch资料，让我们一起努力吧。