PyTorch搭建两层网络

举一个很简单的例子来熟悉使用pytorch构建和训练神经网络的整个过程。

import torch
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)

#定义model
class TowLayerNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        super(TowLayerNet, self).__init__()
        # define the model architecture
        self.linear1 = torch.nn.Linear(D_in, H, bias=False)
        self.linear2 = torch.nn.Linear(H, D_out, bias=False)

    #在调用TowLayerNet(x)的时候自动调用forward函数
    def forward(self, x):
        y_pred = self.linear2(self.linear1(x).clamp(min=0))
        return y_pred

#使用模型
model = TowLayerNet(D_in, H, D_out)

#定义loss函数
loss_fn = nn.MSELoss(reduction='sum')

#定义学习率
learning_rate = 1e-4

#进行梯度下降优化模型
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)

#训练网络
for it in range(500):
    #forward pass
    y_pred = model(x) # model.forward()

    #计算loss
    loss = loss_fn(y_pred, y) # 在计算机中是一个计算图
    print(it, loss.item())

    #每一次要把梯度置0
    optimizer.zero_grad()

    #backward pass,计算出每一步的梯度
    loss.backward()

    #update model paraeters
    optimizer.step()