目录

训练文件：train.py

引入头部文件

 数据转换：

 网络的调用：

损失函数和优化函数：

epoch训练： 

 送入网络：

反向传播更新模型参数：

 保存参数：

训练文件：train.py

引入头部文件

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn

 数据转换：

将其他格式的数据转为Torch类型的数据，举例：

x_train = np.array([[3.3],dtype=np.float32)

y_train = np.array([[1.7],dtype=np.float32)

x_train = torch.from_numpy(x_train)
y_train = torch.from_numpy(y_train)

 网络的调用：

将定义的网络调用出来，有CUDA就将其输入到CUDA中计算，没有就算了

if torch.cuda.is_available():
    model = LinearRegression().cuda()
    print("使用的是： CUDA加速")
else:
    model = LinearRegression()
    print("使用的是： CPU加速")

其中LinearRegression为自定义的网络结构

class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.line = nn.Linear(1,1)  ##输入输出都是一维的

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

损失函数和优化函数：

损失函数和优化函数自行选择，这里示例中用的是均方差，学习率为1e-3，动量加速为0.9

##使用均方误差作为损失函数
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),lr=1e-3,momentum=0.9)

epoch训练： 

自个定义要进行多少轮epoch训练，顺便把数据转变成variable类型，便于后面的求导计算

##进行训练模型
num_epochs = 100
for epoch in range(num_epochs):
    if torch.cuda.is_available():
        inputs = torch.autograd.Variable(x_train).cuda()
        target = torch.autograd.Variable(y_train).cuda()
    else:
        inputs = torch.autograd.Variable(x_train)
        target = torch.autograd.Variable(y_train)

 送入网络：

    ##进行前向传播
    out = model(inputs)
    loss = criterion(out,target)

反向传播更新模型参数：

    ##反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

 保存参数：

后续工作，保存参数或者是保存图片之类的

 

测试

 将需要预测的数据送入模型中得到预测数据，然后进行比对 ：

##预测结果
model.eval()
##这里会报错，需要先将x_train放入cuda中，然后再转为cpu
predict = model(torch.autograd.Variable(x_train.cuda())).cpu()
predict = predict.data.numpy()

plt.plot(x_train.numpy(), predict, 'r.')

plt.show()