Pytorch入門下 —— 其他

WINLSR發表於2021-12-16

本節內容參照小土堆的pytorch入門視訊教程

現有模型使用和修改

pytorch框架提供了很多現有模型,其中torchvision.models包中有很多關於視覺(影像)領域的模型,如下圖:

image-20211214155642948

下面以VGG16為例將講解如何使用以及更改現有模型:

image-20211214161153438

pretrainedTrue,返回在ImageNet上預訓練過的模型;pregressTrue在下載模型時會通過標準錯誤流輸出進度條。

建立如下指令碼並執行:

from torchvision import models


# 建立預訓練過的模型,並輸出進度
vgg16_pretrained = models.vgg16(pretrained=True, progress=True)
# 建立沒訓練過的模型,不輸出進度
vgg16 = models.vgg16(pretrained=False, progress=False)

# 控制檯輸出模型結構
print(vgg16_pretrained)

控制檯輸出如下:

Downloading: "https://download.pytorch.org/models/vgg16-397923af.pth" to C:\Users\winlsr/.cache\torch\hub\checkpoints\vgg16-397923af.pth
100.0%
VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

如上輸出中的的VGG表示模型的class名,featuresVGG含有的一個Sequential元件(Module),avgpoolAdaptiveAvgPool2d元件,classifier同樣為Sequential元件。

建立如下指令碼並執行:

from torchvision import models
from torch import nn


# 建立預訓練過的模型,並輸出進度
vgg16_pretrained = models.vgg16(pretrained=True, progress=True)
# 建立沒訓練過的模型,不輸出進度
vgg16 = models.vgg16(pretrained=False, progress=False)

# 給vgg新增一個線性Module(層)
vgg16_pretrained.add_module("linear", nn.Linear(1000, 10))

# 控制檯輸出模型結構
print(vgg16_pretrained)

輸出如下:

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
  (linear): Linear(in_features=1000, out_features=10, bias=True) # 新增成功
)

建立如下指令碼並執行:

from torchvision import models
from torch import nn


# 建立預訓練過的模型,並輸出進度
vgg16_pretrained = models.vgg16(pretrained=True, progress=True)
# 建立沒訓練過的模型,不輸出進度
vgg16 = models.vgg16(pretrained=False, progress=False)

# 刪除 features 元件
del vgg16_pretrained.features
# 在 classifier 元件中新增元件
vgg16_pretrained.classifier.add_module("7", nn.Linear(1000, 10))
# 修改 classifier 元件中的第1個元件為 softmax(0開始)
vgg16_pretrained.classifier[1] = nn.Softmax()

# 控制檯輸出模型結構
print(vgg16_pretrained)

輸出如下:

VGG(
  # 刪除features成功
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    # 修改為softmax成功
    (1): Softmax(dim=None)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
    # 新增成功
    (7): Linear(in_features=1000, out_features=10, bias=True)
  )
)

模型的儲存與讀取

pytorch中有兩種模型儲存和讀取方式:

執行如下指令碼:

from _07_cifar10_model.cifar10_model import MyModel
import torch

cifar10_model = MyModel()

# 方式1:儲存 模型 + 引數
torch.save(cifar10_model, "cifar10_model.pth")
# 方式2:只儲存 引數(官方推薦)
torch.save(cifar10_model.state_dict(), "cifar10_model_state_dict.pth")

執行成功後,指令碼檔案所在目錄會生成:cifar10_model.pthcifar10_model_state_dict.pth兩個檔案。

恢復方式1儲存的模型:

import torch

# 方式1
cifar10_model = torch.load("cifar10_model.pth")
print(cifar10_model)

輸出如下:

MyModel(
  (model): Sequential(
    (0): Conv2d(3, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (2): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (4): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Flatten(start_dim=1, end_dim=-1)
    (7): Linear(in_features=1024, out_features=64, bias=True)
    (8): Linear(in_features=64, out_features=10, bias=True)
  )
)

恢復方式2儲存的模型(官方推薦):

import torch
from _07_cifar10_model.cifar10_model import MyModel

# 方式2(官方推薦)
cifar10_model = MyModel()
cifar10_model.load_state_dict(torch.load("cifar10_model_state_dict.pth"))
print(cifar10_model)

輸出如下:

MyModel(
  (model): Sequential(
    (0): Conv2d(3, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (2): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (4): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Flatten(start_dim=1, end_dim=-1)
    (7): Linear(in_features=1024, out_features=64, bias=True)
    (8): Linear(in_features=64, out_features=10, bias=True)
  )
)

模型的完整訓練套路

前面我們雖然搭建了在CIFAR10資料集上的分類模型,但是我們並沒有對模型進行完整的訓練。下面會對我們的模型進行一個完整的訓練。訓練程式碼如下:

import time
from torch.utils import tensorboard
from torch.utils.data import DataLoader
from _07_cifar10_model.cifar10_model import MyModel
import torchvision
import torch.nn

if __name__ == "__main__":
    start_time = time.time()
    # 準備訓練資料集和測試資料集
    transform = torchvision.transforms.Compose({
        torchvision.transforms.ToTensor()
    })
    train_data = torchvision.datasets.CIFAR10("./dataset", train=True,
                                              transform=transform,
                                              download=True)
    test_data = torchvision.datasets.CIFAR10("./dataset", train=False,
                                             transform=transform,
                                             download=True)
    train_data_len = len(train_data)
    test_data_len = len(test_data)
    print("訓練集的長度: {}".format(train_data_len))
    print("測試集的長度: {}".format(test_data_len))

    # 建立訓練集和測試集的dataloader
    train_dataloader = DataLoader(dataset=train_data, batch_size=64,
                                  shuffle=True,
                                  num_workers=16)
    test_dataloader = DataLoader(dataset=test_data, batch_size=64,
                                 shuffle=True,
                                 num_workers=16)

    # 建立網路
    cifar10_model = MyModel()

    # 建立損失函式
    loss_func = torch.nn.CrossEntropyLoss()

    # 建立優化器
    # 學習率,科學計數的形式方便改動
    learning_rate = 1e-2
    optimizer = torch.optim.SGD(cifar10_model.parameters(), lr=learning_rate)

    # 訓練次數
    total_train_step = 0
    # 訓練輪次
    epoch = 20

    # 建立 tensorboard SummaryWriter
    writer = tensorboard.SummaryWriter("logs")

    for i in range(epoch):
        print("----------第 {} 輪訓練開始-----------".format(i))

        # 模型進入訓練模式,該方法在當前模型可有可無(加上是個好習慣)
        cifar10_model.train()
        for data in train_dataloader:
            images, targets = data
            outputs = cifar10_model(images)
            loss = loss_func(outputs, targets)

            # 清空上一輪計算的梯度
            optimizer.zero_grad()
            # 反向傳播計算梯度
            loss.backward()
            # 優化器優化引數(執行梯度下降)
            optimizer.step()

            total_train_step += 1
            writer.add_scalar("train/Loss", loss.item(), total_train_step)
            if total_train_step % 100 == 0:
                print("訓練次數: {}, Loss: {}".
                      format(total_train_step, loss.item()))

        total_test_loss = 0.0
        total_accuracy = 0.0
        # 每輪 epoch 後計算模型在測試集上的loss表現
        # 測試時無需計算梯度,可加快計算速度
        # 模型進入驗證(測試)模式,該方法在當前模型可有可無(加上是個好習慣)
        cifar10_model.eval()
        with torch.no_grad():
            for data in test_dataloader:
                images, targets = data
                outputs = cifar10_model(images)
                loss = loss_func(outputs, targets)

                total_test_loss += loss.item()

                accuracy = (outputs.argmax(1) == targets).sum()
                total_accuracy += accuracy
            print("測試準確率:{}".format(total_accuracy/test_data_len))
            writer.add_scalar("test/Loss", total_test_loss, i)
            writer.add_scalar("test/accuracy", total_accuracy/test_data_len, i)

        # 儲存每輪訓練後的模型
        torch.save(cifar10_model.state_dict(),
                   "cifar10_model_state_dict_{}_epoch.pth".format(i))

    writer.close()

    end_time = time.time()
    print("耗時:{}".format(end_time - start_time))

如上程式碼中呼叫模型的train()eval()方法主要是對模型中的DropoutBatchNormModule有用(如果存在),官方解釋如下:

image-20211215122141638 image-20211215122309943

tensorboard視覺化結果如下:

image-20211215142859216

利用GPU訓練

沒有GPU的同學可以想辦法使用google colab,他提供了免費的GPU使用時長,使用起來和jupyter notebook很像。

利用GPU訓練很簡單:

方式一:.cuda()

只需要對 網路模型、資料(輸入、標註)、損失函式呼叫.cuda()方法:

import time
from torch.utils import tensorboard
from torch.utils.data import DataLoader
from _07_cifar10_model.cifar10_model import MyModel
import torchvision
import torch.nn

if __name__ == "__main__":
    start_time = time.time()
    # 準備訓練資料集和測試資料集
    transform = torchvision.transforms.Compose({
        torchvision.transforms.ToTensor()
    })
    train_data = torchvision.datasets.CIFAR10("./dataset", train=True,
                                              transform=transform,
                                              download=True)
    test_data = torchvision.datasets.CIFAR10("./dataset", train=False,
                                             transform=transform,
                                             download=True)
    train_data_len = len(train_data)
    test_data_len = len(test_data)
    print("訓練集的長度: {}".format(train_data_len))
    print("測試集的長度: {}".format(test_data_len))

    # 建立訓練集和測試集的dataloader
    train_dataloader = DataLoader(dataset=train_data, batch_size=64,
                                  shuffle=True,
                                  num_workers=16)
    test_dataloader = DataLoader(dataset=test_data, batch_size=64,
                                 shuffle=True,
                                 num_workers=16)

    # 建立網路
    cifar10_model = MyModel()
    if torch.cuda.is_available():
        cifar10_model = cifar10_model.cuda()

    # 建立損失函式
    loss_func = torch.nn.CrossEntropyLoss()
    if torch.cuda.is_available():
        loss_func = loss_func.cuda()

    # 建立優化器
    # 學習率,科學計數的形式方便改動
    learning_rate = 1e-2
    optimizer = torch.optim.SGD(cifar10_model.parameters(), lr=learning_rate)

    # 訓練次數
    total_train_step = 0
    # 訓練輪次
    epoch = 20

    # 建立 tensorboard SummaryWriter
    writer = tensorboard.SummaryWriter("logs")

    for i in range(epoch):
        print("----------第 {} 輪訓練開始-----------".format(i))

        # 模型進入訓練模式,該方法在當前模型可有可無(加上是個好習慣)
        cifar10_model.train()
        for data in train_dataloader:
            images, targets = data

            if torch.cuda.is_available():
                images = images.cuda()
                targets = targets.cuda()

            outputs = cifar10_model(images)
            loss = loss_func(outputs, targets)

            # 清空上一輪計算的梯度
            optimizer.zero_grad()
            # 反向傳播計算梯度
            loss.backward()
            # 優化器優化引數(執行梯度下降)
            optimizer.step()

            total_train_step += 1
            writer.add_scalar("train/Loss", loss.item(), total_train_step)
            if total_train_step % 100 == 0:
                print("訓練次數: {}, Loss: {}".
                      format(total_train_step, loss.item()))

        total_test_loss = 0.0
        total_accuracy = 0.0
        # 每輪 epoch 後計算模型在測試集上的loss表現
        # 測試時無需計算梯度,可加快計算速度
        # 模型進入驗證(測試)模式,該方法在當前模型可有可無(加上是個好習慣)
        cifar10_model.eval()
        with torch.no_grad():
            for data in test_dataloader:
                images, targets = data

                if torch.cuda.is_available():
                    images = images.cuda()
                    targets = targets.cuda()

                outputs = cifar10_model(images)
                loss = loss_func(outputs, targets)

                total_test_loss += loss.item()

                accuracy = (outputs.argmax(1) == targets).sum()
                total_accuracy += accuracy
            print("測試準確率:{}".format(total_accuracy/test_data_len))
            writer.add_scalar("test/Loss", total_test_loss, i)
            writer.add_scalar("test/accuracy", total_accuracy/test_data_len, i)

        # 儲存每輪訓練後的模型
        torch.save(cifar10_model.state_dict(),
                   "cifar10_model_state_dict_{}_epoch.pth".format(i))

    writer.close()

    end_time = time.time()
    print("耗時:{}".format(end_time - start_time))

方式二:.to()

對 網路模型、資料(輸入、標註)、損失函式呼叫.to()方法,方法中傳入torch.device()物件。這種方式的好處在於不但可以使用GPU,還可以在有多塊GPU時指定使用某塊GPU

如下:

# cpu
CPU_device = torch.device("cpu")
# gpu 只有一塊顯示卡無需指明使用第幾塊
GPU_device = torch.device("cuda")
# 第0塊 gpu
GPU_0_device = torch.device("cuda:0")

完整程式碼如下:

import time
from torch.utils import tensorboard
from torch.utils.data import DataLoader
from _07_cifar10_model.cifar10_model import MyModel
import torchvision
import torch.nn

if __name__ == "__main__":

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    start_time = time.time()
    # 準備訓練資料集和測試資料集
    transform = torchvision.transforms.Compose({
        torchvision.transforms.ToTensor()
    })
    train_data = torchvision.datasets.CIFAR10("./dataset", train=True,
                                              transform=transform,
                                              download=True)
    test_data = torchvision.datasets.CIFAR10("./dataset", train=False,
                                             transform=transform,
                                             download=True)
    train_data_len = len(train_data)
    test_data_len = len(test_data)
    print("訓練集的長度: {}".format(train_data_len))
    print("測試集的長度: {}".format(test_data_len))

    # 建立訓練集和測試集的dataloader
    train_dataloader = DataLoader(dataset=train_data, batch_size=64,
                                  shuffle=True,
                                  num_workers=16)
    test_dataloader = DataLoader(dataset=test_data, batch_size=64,
                                 shuffle=True,
                                 num_workers=16)

    # 建立網路
    cifar10_model = MyModel()
    cifar10_model = cifar10_model.to(device)
    # if torch.cuda.is_available():
    #     cifar10_model = cifar10_model.cuda()

    # 建立損失函式
    loss_func = torch.nn.CrossEntropyLoss()
    loss_func = loss_func.to(device)
    # if torch.cuda.is_available():
    #     loss_func = loss_func.cuda()

    # 建立優化器
    # 學習率,科學計數的形式方便改動
    learning_rate = 1e-2
    optimizer = torch.optim.SGD(cifar10_model.parameters(), lr=learning_rate)

    # 訓練次數
    total_train_step = 0
    # 訓練輪次
    epoch = 20

    # 建立 tensorboard SummaryWriter
    writer = tensorboard.SummaryWriter("logs")

    for i in range(epoch):
        print("----------第 {} 輪訓練開始-----------".format(i))

        # 模型進入訓練模式,該方法在當前模型可有可無(加上是個好習慣)
        cifar10_model.train()
        for data in train_dataloader:
            images, targets = data

            images = images.to(device)
            targets = targets.to(device)
            # if torch.cuda.is_available():
            #     images = images.cuda()
            #     targets = targets.cuda()

            outputs = cifar10_model(images)
            loss = loss_func(outputs, targets)

            # 清空上一輪計算的梯度
            optimizer.zero_grad()
            # 反向傳播計算梯度
            loss.backward()
            # 優化器優化引數(執行梯度下降)
            optimizer.step()

            total_train_step += 1
            writer.add_scalar("train/Loss", loss.item(), total_train_step)
            if total_train_step % 100 == 0:
                print("訓練次數: {}, Loss: {}".
                      format(total_train_step, loss.item()))

        total_test_loss = 0.0
        total_accuracy = 0.0
        # 每輪 epoch 後計算模型在測試集上的loss表現
        # 測試時無需計算梯度,可加快計算速度
        # 模型進入驗證(測試)模式,該方法在當前模型可有可無(加上是個好習慣)
        cifar10_model.eval()
        with torch.no_grad():
            for data in test_dataloader:
                images, targets = data

                images = images.to(device)
                targets = targets.to(device)
                # if torch.cuda.is_available():
                #     images = images.cuda()
                #     targets = targets.cuda()

                outputs = cifar10_model(images)
                loss = loss_func(outputs, targets)

                total_test_loss += loss.item()

                accuracy = (outputs.argmax(1) == targets).sum()
                total_accuracy += accuracy
            print("測試準確率:{}".format(total_accuracy/test_data_len))
            writer.add_scalar("test/Loss", total_test_loss, i)
            writer.add_scalar("test/accuracy", total_accuracy/test_data_len, i)

        # 儲存每輪訓練後的模型
        torch.save(cifar10_model.state_dict(),
                   "cifar10_model_state_dict_{}_epoch.pth".format(i))

    writer.close()

    end_time = time.time()
    print("耗時:{}".format(end_time - start_time))

模型驗證

前面的小節中,我們已經將模型訓練好了,且儲存了每輪訓練後的模型引數。現在我們選擇一個在測試集上表現最好的模型進行恢復,然後在網上隨便找些圖片,看我們的模型能否分類正確。根據tensorboard的顯示,表現最好的模型是在第18輪訓練後的模型,能達到65%左右的正確率。預測圖片如下:

image-20211215160913428 image-20211215163244976

根據CIFAR10資料集中定義,dogtarget5airplanetarget0

image-20211215163104870

預測程式碼如下:

import torch
from PIL import Image
import torchvision
from _07_cifar10_model.cifar10_model import MyModel

dog_img_path = "dog.png"
airplane_img_path = "airplane.png"
dog_img_PIL = Image.open(dog_img_path)
airplane_img_PIL = Image.open(airplane_img_path)
# 將4通道RGBA轉成3通道RGB
dog_img_PIL = dog_img_PIL.convert("RGB")
airplane_img_PIL = airplane_img_PIL.convert("RGB")

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),
    torchvision.transforms.ToTensor()
])
dog_img_tensor = transform(dog_img_PIL)
airplane_img_tensor = transform(airplane_img_PIL)
# print(dog_img_tensor.shape)
dog_img_tensor = torch.reshape(dog_img_tensor, (-1, 3, 32, 32))
airplane_img_tensor = torch.reshape(airplane_img_tensor, (1, 3, 32, 32))

cifar10_model = MyModel()
cifar10_model.load_state_dict(torch.load(
    "../_10_train_model/cifar10_model_state_dict_18_epoch.pth"))

cifar10_model.eval()
with torch.no_grad():
    output = cifar10_model(dog_img_tensor)
    print(output.argmax(1))
    output = cifar10_model(airplane_img_tensor)
    print(output.argmax(1))

輸出如下:

tensor([7]) # 預測錯誤
tensor([0]) # 預測正確

相關文章