比如,建了一个两层全连接层的神经网络:
class LinearClassifier_2layer(nn.Module): def __init__(self, last_layer_dim=None, n_label=None): super(LinearClassifier_2layer, self).__init__() self.classifier = nn.Sequential( nn.Linear(last_layer_dim, 2*last_layer_dim), nn.ReLU(), nn.Linear(2*last_layer_dim, n_label) ) def forward(self, x): return self.classifier(x) model = LinearClassifier_2layer(512, 10)
然后,用这个神经网络训练。训练过程中,如果想查看权重或梯度:
model.classifier[0].weight[:3, :3] # 查看第一层全连接层的一部分权重 model.classifier[0].weight.grad[:3, :3] # 查看第一层全连接层的一部分梯度 model.classifier[1].weight[:3, :3] # 查看第二层全连接层的一部分权重 model.classifier[1].weight.grad[:3, :3] # 查看第一层全连接层的一部分梯度