比如,建了一個兩層全連接層的神經網絡:
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] # 查看第一層全連接層的一部分梯度
