torch_02_多项式回归

"""
torch.float64对应torch.DoubleTensor
torch.float32对应torch.FloatTensor
将真实函数的数据点能够拟合成一个多项式
eg：y = 0.9 +0.5×x + 3×x*x + 2.4 ×x*x*x
"""
import torch

from torch import nn

def make_features(x):
    x = x.unsqueeze(1)#在原来的基础上扩充了一维
    return torch.cat([x ** i for i in range(1,4)], 1)


def get_batch(batch_size=32):

    random = torch.randn(batch_size)
    # print('random')
    # print(random) #32个数

    x = make_features(random)#进行维度扩充，扩充后32*1，又进行1,2,3次幂运算，拼接后32*3

    '''Compute the actual results'''
    y = f(x) # 32*3 *3*1
    if torch.cuda.is_available():
        return torch.autograd.Variable(x).cuda(), torch.autograd.Variable(y).cuda()
    else:
        return torch.autograd.Variable(x), torch.autograd.Variable(y)


w_target = torch.FloatTensor([0.5,3,2.4]).unsqueeze(1)#三行一列
b_target = torch.FloatTensor([0.9])


def f(x):
    return x.mm(w_target)+b_target[0]

class poly_model(nn.Module):
    def __init__(self):
        super(poly_model, self).__init__()
        self.poly = nn.Linear(3, 1)# 输入是3维，输出是1维

    def forward(self, x):
        out = self.poly(x)
        return out

if torch.cuda.is_available():
 model = poly_model().cuda()
else:
 model = poly_model()

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

epoch = 0
for epoch in range(20):
    batch_x,batch_y = get_batch()#batch_x 和get_batch里面的x是一样的
    output = model(batch_x)
    loss = criterion(output,batch_y)
    print_loss = loss
    print(loss.item()) # 0.4版本之后使用loss.item()从标量中获得Python number
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
print('finished')