深度学习之PyTorch---- 多项式线性回归

"""
多项式回归
"""
def make_features(x):
    """Builds 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)


def f(x):
    """Approximated function."""
    return x.mm(W_target) + b_target[0]

def get_batch(batch_size=100):
    """Builds a batch (x , f(x)) pair"""
    random = torch.randn(batch_size)
    random = np.sort(random)
    random = torch.Tensor(random)
    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)
        
    def forward(self,x):
        out = self.poly(x)
        return out
    


if __name__ == '__main__':
    epoch = 0
    W_target = torch.FloatTensor([0.5,3,2.4]).unsqueeze(1)
    b_target = torch.FloatTensor([0.9])
    model = poly_model()
    
    # 定义损失函数和优化器
    criterion = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(),lr=1e-3)
    
    while True:
        # Get Data
        batch_x ,batch_y = get_batch()

        # Forward pass
        ouput = model(batch_x)
        loss = criterion(ouput,batch_y)
        print_loss = loss.data[0]

        # Reset gradient
        optimizer.zero_grad()

        # Backward pass
        loss.backward()
        optimizer.step()
        epoch+=1
        if print_loss < 1e-3:
            break

    print('Loss {:.6f} after {} batches'.format(loss.item(),epoch))
    print('==> Learned function:  y = {:.2f} + {:.2f}*x + {:.2f}*x^2 + {:.2f}*x^3'.format(model.poly.bias[0],
                                                                                         model.poly.weight[0][0],
                                                                                         model.poly.weight[0][1],
                                                                                         model.poly.weight[0][2]))
    print('==> Actual function:   y = {:.2f} + {:.2f}*x + {:.2f}*x^2 + {:.2f}*x^3'.format(b_target[0],
                                                                                          W_target[0][0],
                                                                                          W_target[1][0],
                                                                                          W_target[2][0]))

    # Prediction

    predict = model(batch_x)

    x = batch_x.numpy()[:,0]
    plt.plot(x,batch_y.numpy(),'ro',label='real cure')

    predict = predict.data.numpy()
    plt.plot(x,predict,'b',label='fitting cure')
    plt.legend()

输出：

Loss 0.000998 after 1426 batches
==> Learned function:  y = 0.94 + 0.53*x + 2.98*x^2 + 2.39*x^3
==> Actual function:   y = 0.90 + 0.50*x + 3.00*x^2 + 2.40*x^3