小程序

小程序

import torch
import matplotlib.pyplot as plt

x_data = torch.Tensor([[1.0],[2.0],[3.0],[4.0],[5.0]])
y_data = torch.Tensor([[2.0],[4.0],[6.0],[8.0],[10.0]])

class MODEL(torch.nn.Module):
    def __init__(self,n_input,n_output):
        super(MODEL, self).__init__()
        self.linear = torch.nn.Linear(n_input,n_output)

    def forward(self, x):
        out = self.linear(x)
        return out
model = MODEL(n_input=1,n_output=1)

optimizer = torch.optim.SGD(model.parameters(),lr=0.01)
loss_fuc = torch.nn.MSELoss()

epoch_list = []
loss_list = []


for epoch in range(100):
    outputs = model(x_data)
    loss = loss_fuc(outputs,y_data)

    epoch_list.append(epoch)
    loss_list.append(loss)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    print('epoch:%d'%epoch,'loss:%f'%loss.item())

print('w=',model.linear.weight.item())
print('b=',model.linear.bias.item())
x_test = torch.Tensor([6.0])
y_test = model(x_test)
print('y_pre=',y_test.data.numpy())

plt.plot(epoch_list,loss_list)
plt.title('linearmodel')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.show()

参考：https://blog.csdn.net/qq_41521512/article/details/106646318?utm_medium=distribute.pc_relevant.none-task-blog-OPENSEARCH-1.control&dist_request_id=&depth_1-utm_source=distribute.pc_relevant.none-task-blog-OPENSEARCH-1.control