1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100
| import torch.nn as nn import torch.optim as optim import torch.nn.functional as F
class TheModelClass(nn.Module): def __init__(self): super(TheModelClass,self).__init__() self.conv1=nn.Conv2d(3, 6, 5) self.pool=nn.MaxPool2d(2, 2) self.conv2=nn.Conv2d(6, 16, 5) self.fc1=nn.Linear(16*5*5, 120) self.fc2=nn.Linear(120, 84) self.fc3=nn.Linear(84, 10)
def forward(self,x): x=self.pool(F.relu(self.conv1(x))) x=self.pool(F.relu(self.conv2(x))) x=x.view(-1,16*5*5) x=F.relu(self.fc1(x)) x=F.relu(self.fc2(x)) x=self.fc3(x) return x
def main(): model = TheModelClass()
optimizer=optim.SGD(model.parameters(), lr=0.001, momentum=0.9) ''' model的state_dict()与optimizer的略有不同 model: torch.nn.Module模块中的state_dict只包含卷积层和全连接层的参数 当网络中存在batchnorm时,例如vgg网络结构,torch.nn.Module模块中 的state_dict也会存放batchnorm's running_mean optimizer: state_dict字典对象包含state和param_groups的字典对象,而param_groups key 对应的value也是一个由学习率,动量等参数组成的一个字典对象 ''' print('Model.state_dict: ') model_param = model.state_dict() for param_tensor in model_param: print(param_tensor, '\t', model.state_dict()[param_tensor].size()) ''' Model.state_dict: conv1.weight torch.Size([6, 3, 5, 5]) conv1.bias torch.Size([6]) conv2.weight torch.Size([16, 6, 5, 5]) conv2.bias torch.Size([16]) fc1.weight torch.Size([120, 400]) fc1.bias torch.Size([120]) fc2.weight torch.Size([84, 120]) fc2.bias torch.Size([84]) fc3.weight torch.Size([10, 84]) fc3.bias torch.Size([10]) '''
print('Optimizer state_dict: ') optim_param = optimizer.state_dict() for var_name in optim_param: print(var_name, '\t', optimizer.state_dict()[var_name]) ''' Optimizer state_dict: state {} param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]}] '''
params_groups = optimizer.param_groups print("optimizer.param_groups: ") print(params_groups) ''' 此处会打印出每一个参数的具体数值 这里假设是torch.nn.Linear(1, 1),有两个参数 w&&b [{'params': [Parameter containing: tensor([[0.8117]], device='cuda:0', requires_grad=True), Parameter containing: tensor([0.8024], device='cuda:0', requires_grad=True)], 'lr': 0.001, 'momentum': 0, 'dampening': 0, 'weight_decay': 0, 'nesterov': False}] '''
if __name__=='__main__': main()
|