pytorch构建两层神经网络并使用自动求导机制

'''
使用pytorch的自动求导机制
'''
import torch

dtype=torch.float
device=torch.device("cpu")

# N是批量大小; D_in是输入维度;
# H是隐藏的维度; D_out是输出维度。

N,D_in,H,D_out=64,1000,100,10

#随机构建输入输出
x=torch.randn(N,D_in,device=device,dtype=dtype)
y=torch.randn(N,D_out,device=device,dtype=dtype)

w1=torch.randn(D_in,H,device=device,dtype=dtype,requires_grad=True)
w2=torch.randn(H,D_out,device=device,dtype=dtype,requires_grad=True)

#设置学习率
learning_rate=1e-6
for t in range(500):
    y_pred=x.mm(w1).clamp(min=0).mm(w2)#前馈计算
    loss = (y_pred - y).pow(2).sum()
    print(t,loss.item())
    loss.backward()#自动求导

    with torch.no_grad():#更新阶段不添加计算图，去除求导机制
        w1-=learning_rate*w1.grad
        w2-=learning_rate*w2.grad
        #手动将梯度设为0
        w1.grad.zero_()
        w2.grad.zero_()