pytorch实现两层BP网络

import torch.nn as nn

N, D_in, H, D_out = 64, 1000, 100, 10

# 随机创建一些训练数据
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

# 定义两层BP网络的类
class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        super(TwoLayerNet, self).__init__()
        # define the model architecture
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)
    
    def forward(self, x):
        y_pred = self.linear2(self.linear1(x).clamp(min=0))  #torch.nn.clamp类似与numpy的np.maximum(relu函数的作用)
        return y_pred

model = TwoLayerNet(D_in, H, D_out)
loss_fn = nn.MSELoss(reduction='sum')   #选择MSE作为损失loss
learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for it in range(5):
    # Forward pass，前向网络
    y_pred = model(x) # model.forward() 
    
    # compute loss，损失函数
    loss = loss_fn(y_pred, y) # computation graph
    print(it, loss.item())

    optimizer.zero_grad()   #梯度归零
    # Backward pass，求导
    loss.backward()
    
    # update model parameters，参数更新
    optimizer.step()

out

0 678.2354736328125
1 660.7021484375
2 643.60888671875
3 627.0241088867188
4 610.9411010742188