Pytorch搭建神经网络入门

简介

Pytorch作为当前最受欢迎的深度学习框架之一，是笔者十分推荐大家学习和使用的。以下代码，我使用Pytorch提供的线性函数简单搭建了一个神经网络，用以实现对一次函数y = 2*x + 1的拟合，用以帮助大家实现从 0 到 1 搭建神经网络模型。

代码

以下为详细代码：

import torch
import torch.nn as nn
import torch.optim as optim

import numpy as np


# 搭建神经网络模型
class Model(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Model, self).__init__()
        self.layer = nn.Sequential(
            nn.Linear(in_features=input_dim, out_features=output_dim)
        )

    def forward(self, x):
        out = self.layer(x)
        return out


if __name__ == '__main__':
    # 训练数据处理
    train_set = []
    for i in range(100):
        train_set.append([i, 2 * i + 1])

    x_train = np.array([[s[0]] for s in train_set], dtype=np.float32)
    y_train = np.array([[s[1]] for s in train_set], dtype=np.float32)

    # 训练部分
    epoches = 5000
    learning_rate = 0.01

    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    model = Model(1, 1).to(device)

    # 定义损失函数
    criterion = nn.MSELoss()
    # 定义优化器
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    for epoch in range(epoches):
        inputs = torch.from_numpy(x_train).to(device)
        targets = torch.from_numpy(y_train).to(device)

        # 梯度清零，否则将对上一次的迭代结果进行累加
        optimizer.zero_grad()
        
        # 根据 inputs 和 model 计算出outputs
        outputs = model(inputs)
        
        # 计算 loss 值并进行后向传播
        loss = criterion(outputs, targets)
        loss.backward()
        
        # 根据优化器对参数进行优化
        optimizer.step()
        
        # 以 50 个epoch为单位输出当前 loss 值
        if epoch % 50 == 0:
            print('epoch:{}, loss: {}'.format(epoch, loss.item()))

    # 保存模型，此处保存模型的状态字典
    torch.save(model.state_dict(), 'linear.pt')

    # 实例化模型并加载之前训练好的模型参数
    pre_model = Model(1, 1)
    pre_model.load_state_dict(torch.load('linear.pt'))

    # 预测数据
    predicted = pre_model(torch.from_numpy(np.array([[200], [201]], dtype=np.float32)).requires_grad_()).data.numpy()
    print(predicted)

问题探讨

待更新…