基于PyTorch使用Jupyter Notebook环境实现线性回归算法; 使用均方差作为损失函数; 使用随机批梯度下降作为参数优化算法; 动画显示模型拟合数据的过程; 序列化模型参数;
from __future__ import print_function
import torch
import torch. nn as nn
import numpy as np
import matplotlib. pyplot as plt
import time
% matplotlib inline
from IPython import display
input_size = 1
output_size = 1
num_epochs = 60
learning_rate = 0.001
x_train = np. array( [ [ 3.3 ] , [ 4.4 ] , [ 5.5 ] , [ 6.71 ] , [ 6.93 ] , [ 4.168 ] ,
[ 9.779 ] , [ 6.182 ] , [ 7.59 ] , [ 2.167 ] , [ 7.042 ] ,
[ 10.791 ] , [ 5.313 ] , [ 7.997 ] , [ 3.1 ] ] , dtype= np. float32)
y_train = np. array( [ [ 1.7 ] , [ 2.76 ] , [ 2.09 ] , [ 3.19 ] , [ 1.694 ] , [ 1.573 ] ,
[ 3.366 ] , [ 2.596 ] , [ 2.53 ] , [ 1.221 ] , [ 2.827 ] ,
[ 3.465 ] , [ 1.65 ] , [ 2.904 ] , [ 1.3 ] ] , dtype= np. float32)
model = nn. Linear( input_size, output_size)
criterion = nn. MSELoss( )
optimizer = torch. optim. SGD( model. parameters( ) , lr= learning_rate)
for epoch in range ( num_epochs) :
inputs = torch. from_numpy( x_train)
targets = torch. from_numpy( y_train)
outputs = model( inputs)
loss = criterion( outputs, targets)
optimizer. zero_grad( )
loss. backward( )
optimizer. step( )
if ( epoch+ 1 ) % 5 == 0 :
predicted = model( torch. from_numpy( x_train) ) . detach( ) . numpy( )
plt. plot( x_train, y_train, 'ro' , label= 'Original data' )
plt. plot( x_train, predicted, label= 'Fitted line' )
plt. xlim( ( 0 , 12 ) )
plt. ylim( ( - 2 , 8 ) )
plt. legend( )
plt. show( )
display. clear_output( wait= True )
plt. pause( 1 )
predicted = model( torch. from_numpy( x_train) ) . detach( ) . numpy( )
plt. plot( x_train, y_train, 'ro' , label= 'Original data' )
plt. plot( x_train, predicted, label= 'Fitted line' )
plt. legend( )
plt. show( )
torch. save( model. state_dict( ) , 'model.ckpt' )