PyTorch基礎-Adam優化器使用-06
阿新 • • 發佈:2021-02-10
技術標籤:Pytorch深度學習神經網路機器學習pytorchadam
當不知道使用什麼優化器的時候可以使用adam優化器
程式碼
import numpy as np
import torch
from torch import nn,optim
from torch.autograd import Variable
from torchvision import datasets,transforms
from torch.utils.data import DataLoader
# 訓練集
train_data = datasets.MNIST(root="./" , # 存放位置
train = True, # 載入訓練集
transform=transforms.ToTensor(), # 把資料變成tensor型別
download = True # 下載
)
# 測試集
test_data = datasets.MNIST(root="./",
train = False,
transform=transforms.ToTensor(),
download = True
)
# 批次大小
batch_size = 64
# 裝載訓練集
train_loader = DataLoader(dataset=train_data,batch_size=batch_size,shuffle=True)
# 裝載測試集
test_loader = DataLoader(dataset=test_data, batch_size=batch_size,shuffle=True)
for i,data in enumerate(train_loader):
inputs,labels = data
print(inputs.shape)
print(labels.shape)
break
# 定義網路結構
class Net(nn.Module):
def __init__(self):
super(Net,self).__init__()# 初始化
self.fc1 = nn.Linear(784,10) # 784個輸入10個輸出
self.softmax = nn.Softmax(dim=1) # 啟用函式 dim=1表示對第一個維度進行概率計算
def forward(self,x):
# torch.Size([64, 1, 28, 28]) -> (64,784)
x = x.view(x.size()[0],-1) # 4維變2維 (在全連線層做計算只能2維)
x = self.fc1(x) # 傳給全連線層繼續計算
x = self.softmax(x) # 使用softmax啟用函式進行計算
return x
# 定義模型
model = Net()
# 定義代價函式
mse_loss = nn.CrossEntropyLoss()# 交叉熵
# 定義優化器
optimizer = optim.Adam(model.parameters(),lr=0.001)# Adam梯度下降
# 定義模型訓練和測試的方法
def train():
for i,data in enumerate(train_loader):
# 獲得一個批次的資料和標籤
inputs,labels = data
# 獲得模型預測結果(64,10)
out = model(inputs)
# 交叉熵代價函式out(batch,C:類別的數量),labels(batch)
loss = mse_loss(out,labels)
# 梯度清零
optimizer.zero_grad()
# 計算梯度
loss.backward()
# 修改權值
optimizer.step()
def test():
correct = 0
for i,data in enumerate(test_loader):
# 獲得一個批次的資料和標籤
inputs,labels = data
# 獲得模型預測結果(64,10)
out = model(inputs)
# 獲得最大值,以及最大值所在的位置
_,predicted = torch.max(out,1)
# 預測正確的數量
correct += (predicted==labels).sum()
print("Test acc:{0}".format(correct.item()/len(test_data)))
# 訓練
for epoch in range(10):
print("epoch:",epoch)
train()
test()