1. 程式人生 > 其它 >PyTorch基礎-Adam優化器使用-06

PyTorch基礎-Adam優化器使用-06

技術標籤: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()

在這裡插入圖片描述