Pytorch学习（2）多项式回归实现

import torch
from torch import nn
from torch.autograd import Variable
from torch import optim
import numpy as np
import matplotlib.pyplot as plt
 
def make_features(x):
    '''Build features a matrix with columns [x, x^2, x^3]. '''
    x = x.unsqueeze(1)
    return torch.cat([x ** i for i in range(1, 4)], 1)
 
W_target = torch.FloatTensor([0.5, 3, 2.4]).unsqueeze(1)
b_target = torch.FloatTensor([0.9])
 
def f(x):
    '''Approximated function '''
    return x.mm(W_target) + b_target[0]
 
def get_batch(batch_size = 32):
    '''Buile a batch pair'''
    random = torch.randn(batch_size)
    x = make_features(random)
    y = f(x)
    return Variable(x), Variable(y)
 
# define model
class poly_model(nn.Module):
    def __init__(self):
        super(poly_model, self).__init__()
        self.poly = nn.Linear(3, 1)   # three inputs and one output
        
    def forward(self, x):
        out = self.poly(x)
        return out 

# define model and loss  
model = poly_model()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr = 1e-3)
 
epoch = 0
loss_all = []
while True:
    #get data
    batch_x, batch_y = get_batch()
    
    #forward pass
    output = model(batch_x)
    loss = criterion(output, batch_y)
    print_loss = loss.data.item()
    
    #reset gradients
    optimizer.zero_grad()
    
    #backward pass
    loss.backward()
    
    #update parameters
    optimizer.step()
    
    epoch += 1
    
    print('epoch {}'.format(epoch + 1))
    print('loss : {}'.format(print_loss))
    loss_all.append(print_loss)
    
    if print_loss < 1e-3:
        break

# plot line
plt.figure()
index = np.arange(len(loss_all))
plt.plot(index, loss_all, 'r')
plt.legend(['train_loss'])
plt.show()

# show result
#Actual function
x = np.arange(-1, 1, 0.1)
w0 = W_target[0].item()
w1 = W_target[1].item()
w2 = W_target[2].item()
b = b_target.item()
y = w0 * x + w1 * (x ** 2) + w2 * (x ** 3) + b

#
W_predict0, W_predict1, W_predict2 = model.poly.weight[0]
w_0 = W_predict0.data.item()
w_1 = W_predict1.data.item()
w_2 = W_predict2.data.item()
b_ = model.poly.bias[0].item()
y_ = w_0 * x + w_1 * (x ** 2) + w_2 * (x ** 3) + b_

plt.figure()
plt.plot(x, y, 'ro')
plt.plot(x, y_, 'b-')
plt.legend(['actual curve', 'predict curve'])
plt.show()

问题探究：

（1）unsqueeze( )和sequeeze( )函数

（2）程序还有需要完善的部分