PyTorch搭建神经网络的一般步骤

  • 给定输入输出
  • 定义一个模型
  • 定义损失函数(loss function)和优化函数(optimizer)
  • 训练一个过程

一个简单的例子:

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)

class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        super(TwoLayerNet, self).__init__()
        self.linear1 = nn.Linear(D_in, H, bias=False)
        self.linear2 = nn.Linear(H, D_out, bias=False)
        
    def forward(self, x):
        y_pred = self.linear2(self.linear1(x).clamp(min=0))
        return y_pred
    
    
# model = TwoLayerNet(D_in, H, D_out)

if torch.cuda.is_available():
    x = x.cuda()
    y = y.cuda()
    model = TwoLayerNet(D_in, H, D_out).cuda()

else:
    model = TwoLayerNet(D_in,H, D_out)
    
# 损失函数和优化函数
loss_fn = nn.MSELoss(reduction='sum')

optimizer = torch.optim.SGD(model.parameters(), lr = 1e-4)

for it in range(500):
    # Forward pass
    y_pred = model(x)
    
    # 计算loss
    loss = loss_fn(y_pred, y)
    print(it, loss.item())
    
    # backward pass
    optimizer.zero_grad()  # 避免出现梯度累积
    loss.backward()   # 反向传播
    
    # updata model parameters
    optimizer.step()
    
    


训练结果:

PyTorch搭建神经网络的一般步骤_第1张图片

你可能感兴趣的:(深度学习入门)