自留,面试时手撕用
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)