用pytorch实现神经网络线性回归
import torch
import matplotlib.pyplot as plt
from torch import nn, optim
from time import perf_counter
x = torch.unsqueeze(torch.linspace(-3, 3, 10000), dim=1)
y = x + 1.2 * torch.rand(x.size())
class LR(nn.Module):
def __init__(self):
super(LR, self).__init__()
self.linear = nn.Linear(1, 1)
def forward(self, x):
out = self.linear(x)
return out
CUDA = torch.cuda.is_available()
if CUDA:
LR_model = LR().cuda()
inputs = x.cuda()
target = y.cuda()
else:
LR_model = LR()
inputs = x
target = y
criterion = nn.MSELoss()
optimizer = optim.SGD(LR_model.parameters(), lr=1e-4)
def draw(output, loss):
if CUDA:
output = output.cpu()
plt.cla()
plt.scatter(x.numpy(), y.numpy())
plt.plot(x.numpy(), output.data.numpy(), 'r-', lw=5)
plt.text(0.5, 0, 'loss=%s'%(loss.item()), fontdict={'size':'20', 'color':'red'})
plt.pause(0.005)
def train(model, criterion, optimizer, epochs):
for epoch in range(epochs):
output = model(inputs)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 80 == 0:
draw(output, loss)
return model, loss
start = perf_counter()
LR_model, loss = train(LR_model, criterion, optimizer, 20000)
finish = perf_counter()
time = finish - start
print("计算时间:%s" % time)
print("final loss:", loss.item())
print("weights:", list(LR_model.parameters()))