本節內容參照小土堆的pytorch
入門視訊教程。
現有模型使用和修改
pytorch
框架提供了很多現有模型,其中torchvision.models
包中有很多關於視覺(影像)領域的模型,如下圖:
下面以VGG16
為例將講解如何使用以及更改現有模型:
pretrained
為True
,返回在ImageNet
上預訓練過的模型;pregress
為True
在下載模型時會通過標準錯誤流輸出進度條。
建立如下指令碼並執行:
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
名,features
是VGG
含有的一個Sequential
元件(Module
),avgpool
是AdaptiveAvgPool2d
元件,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.pth
、cifar10_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()
方法主要是對模型中的Dropout
、BatchNorm
等Module
有用(如果存在),官方解釋如下:
tensorboard
視覺化結果如下:
利用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%左右的正確率。預測圖片如下:
根據CIFAR10
資料集中定義,dog
的target
為5
,airplane
的target
為0
:
預測程式碼如下:
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]) # 預測正確