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

最新推荐文章于 2025-01-14 20:04:15 发布

原创

最新推荐文章于 2025-01-14 20:04:15 发布 · 978 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#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 = m