pytorch几十行代码搞清楚模型的构建和训练

import torch
import torch.nn as nn

N, D_in, H, D_out = 64, 1000, 100, 10
# data
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

# mdoel define
class TwoLayerNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        # main layers
        super(TwoLayerNet, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.linear2 = nn.Linear(H, D_out)
        
    def forward(self, x):
        y_pred = self.linear2(self.linear1(x).clamp(min=0))
        return y_pred
    
# init model
loss_fn = nn.MSELoss(reduction='sum')
model = TwoLayerNet(D_in, H, D_out)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# training
for i in range(500):
    # 1.forward pass
    y_pred = model(x)
    
    # 2.compute loss
    loss = loss_fn(y_pred, y)
    print(i, loss.item())
    
    optimizer.zero_grad()
    # 3.backward pass
    loss.backward()
    
    # 4.weights update
    optimizer.step()