pytorch 构建神经网络流程

import torch
import torch.nn as nn

dtype = torch.float
device = torch.device("cpu")


# init 定义网络层，forward 定义如何前向传播（x * layer * activate） return loss 或者 return y_pred
class TwoLayerNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        super(TwoLayerNet, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.linear2 = nn.Linear(H, D_out)

    def forward(self, x):
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred


# 定义数据，和神经元个数等
N, D_in, H, D_out = 64, 1000, 100, 10

x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

# 定义实例 model
model = TwoLayerNet(D_in, H, D_out)

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

# 定义优化器
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 迭代
for t in range(500):
    # 模型预测值
    y_pred = model(x)

    # 计算 loss
    loss = criterion(y_pred, y)

    # 梯度初始化为零
    optimizer.zero_grad()

    # 反向传播，计算 loss
    loss.backward()

    # 反向传播求梯度
    optimizer.step()

# model
torch.save(model, './model/xxx.pth')  # 直接保存模型
# 保存模型参数
torch.save(model.state_dict(), './model/xxxxx.pth')

预测的时候需要先定义网络：

import torch
import torch.nn as nn

dtype = torch.float
device = torch.device("cpu")


# init 定义网络层，forward 定义如何前向传播（x * layer * activate） return loss 或者 return y_pred
class TwoLayerNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        super(TwoLayerNet, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.linear2 = nn.Linear(H, D_out)

    def forward(self, x):
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred


# 定义数据，和神经元个数等
N, D_in, H, D_out = 64, 1000, 100, 10

x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

# 定义实例 model
model = TwoLayerNet(D_in, H, D_out)

# 模型类必须在此之前被定义，两种方式 load 模型
net = torch.load('./model/xxx.pth')
net.load_state_dict(torch.load('./model/xxxxx.pth'))

# 运行推理之前，务必调用 model.eval() 设置 dropout 和 batch normalization 层为评估模式。如果不这么做，可能导致模型推断结果不一致。
net.eval()

# 预测不是用 predict 方法，而是： 模型(特征tensor) 
prediction = net(feature_x)