究极简单pytorch流程

自留，面试时手撕用

import torch
import numpy as np
import torch.nn as nn

class Net(nn.Module):

    def __init__(self, features):
        super(Net, self).__init__()

        self.linear1 = nn.Linear(features, 128)
        self.linear2 = nn.Linear(128, 256)
        self.linear3 = nn.Linear(256, 1)
        self.relu = nn.ReLU()

    def forward(self, x):
        y_pred = self.linear1(x)
        y_pred = self.relu(y_pred)
        y_pred = self.linear2(y_pred)
        y_pred = self.relu(y_pred)
        y_pred = self.linear3(y_pred)
        return y_pred


train_features = torch.tensor(np.random.rand(100,567),dtype=torch.float) # print train_features.shap
train_labels = torch.tensor(np.random.rand(100,1),dtype=torch.float).view(-1, 1)
test_features = torch.tensor(np.random.rand(30,567),dtype=torch.float)

model = Net(features=train_features.shape[1])
criterion = nn.MSELoss(reduction='mean')
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

losses = []
epoch = 500
for t in range(epoch):
    y_pred = model(train_features)
    loss = criterion(y_pred, train_labels)
    print(t, loss.item())
    losses.append(loss.item())
    if torch.isnan(loss):
        break
    # PyTorch的backward()方法计算梯度会默认将本次计算的梯度与缓存中已有的梯度加和。
    optimizer.zero_grad() # 将模型中各参数的梯度清零, 必须在反向传播前先清零
    loss.backward() # 反向传播，计算各参数对于损失loss的梯度
    optimizer.step() # 根据刚刚反向传播得到的梯度更新模型参数

predictions = model(test_features).detach().numpy()
print(predictions)