pytorch（2）多项式回归

"""多项式回归代码实现"""

import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
def make_features(x):
    """Builds features i.e. a matrix with columns [x, x^2, x^3]."""
    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=50):
    """Builds a batch i.e. (x, f(x)) pair."""
    x = torch.unsqueeze(torch.linspace(-3, 3, 100), dim=1)

    x_columns = make_features(x)
    y = f(x_columns)
    if torch.cuda.is_available():
        
        return Variable(x_columns).cuda(), Variable(y).cuda(),x
    else:
        
        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__':

    W_target = torch.FloatTensor([0.5, 3, 2.4]).unsqueeze(1)#定各项系数，y=0.5x^3+3x^2+2.4x+0.9
    b_target = torch.FloatTensor([0.9])#定常数
    if torch.cuda.is_available():
        model = poly_model().cuda()
    else:
        model = poly_model()

    criterion = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=1e-3)

    epoch = 0
    while True:
        # Get data
        inputs, target,x = get_batch()

        # Forward pass
        output = model(inputs)
        loss = criterion(output, target)
        print_loss = loss.item()
        # Reset gradients
        optimizer.zero_grad()
        # Backward pass
        loss.backward()
        # update parameters
        optimizer.step()
        epoch += 1
        
        if print_loss < 1e-3:
            break
        
        if epoch % 50 == 0:
        # plot and show learning process
            plt.cla()
            inputs = inputs.cpu()
            target = target.cpu()
            output = output.cpu()
            loss = loss.cpu()
            plt.scatter(x.data.numpy(), target.data.numpy())
            plt.plot(x.data.numpy(), output.data.numpy(), 'r-', lw=5)
            plt.text(0.5, 0, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color':  'red'})
            plt.pause(0.1)
    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]))

            
    plt.ioff()
    plt.show()