import torchimport numpy as npimport random
class PolyRegression(torch.nn.Module): def __init__(self, input_dim): super(PolyRegression, self).__init__() self.linear = torch.nn.Linear(in_features=input_dim, out_features=1)
def forward(self, x): return self.linear(x)
def generate_data(batch_size=32): # 产生正态分布随机数、添加一个维度 # 32行1列 x = torch.randn(batch_size).unsqueeze(1) # f(x) = 2 * x^3 + 3 * x^2 + 4 * x + 5 # 32行1列 y = 2.0 * x**3 + 3.0 * x**2 + 4.0 * x + 5.0 + random.random()/10
# 将x, x平方、x三次方作为输入 # 按维度1拼接（横着拼） # 32行3列 x_data = torch.cat([x**i for i in range(1, 4)], 1) return (x_data, y)
def train(epoch=1000): device = 'cuda' if torch.cuda.is_available() else 'cpu' model = PolyRegression(3).to(device) loss_func = torch.nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.01) for i in range(epoch): x, y = generate_data(32) x = x.to(device) y = y.to(device)
# PyTorch标准流程 predicted = model(x) loss = loss_func(predicted, y) # 梯度清零 optimizer.zero_grad() # 反向传播 loss.backward() # 优化参数 optimizer.step()
print(f'Epoch {i+1}, loss {loss.item()}')
# 最终结果 print(f'Weight: {model.linear.weight.data}') print(f'Bias: {model.linear.bias.data}')
if __name__ == '__main__': train(2000)
评论