訓練模型並驗證(CIFAR10數據集)

CIFAR10這個數據集我們之前已經使用過了。是一個擁有十個分類的模型。本期我們進行一次完整的模型並測試。

import torchimport torchvisionfrom torch.utils.tensorboard import SummaryWriter
from torch import nnfrom torch.utils.data import DataLoader
device = torch.device("cuda")
train_data = torchvision.datasets.CIFAR10(root="./dataset_CIFAR10", train=True, transform=torchvision.transforms.ToTensor(),                                          download=True)test_data = torchvision.datasets.CIFAR10(root="./dataset_CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),                                         download=True)
train_data_size = len(train_data)test_data_size = len(test_data)print("訓練數據集的長度為：{}".format(train_data_size))print("測試數據集的長度為：{}".format(test_data_size))

train_dataloader = DataLoader(train_data, batch_size=64)test_dataloader = DataLoader(test_data, batch_size=64)
class Huaji(nn.Module):    def __init__(self):        super(Huaji, self).__init__()        self.model = nn.Sequential(            nn.Conv2d(3, 32, 5, 1, 2),            nn.MaxPool2d(2),            nn.Conv2d(32, 32, 5, 1, 2),            nn.MaxPool2d(2),            nn.Conv2d(32, 64, 5, 1, 2),            nn.MaxPool2d(2),            nn.Flatten(),            nn.Linear(64*4*4, 64),            nn.Linear(64, 10)        )
    def forward(self, x):        x = self.model(x)        return xhuaji = Huaji()huaji = huaji.to(device)
loss_fn = nn.CrossEntropyLoss()loss_fn = loss_fn.to(device)learning_rate = 1e-2optimizer = torch.optim.SGD(huaji.parameters(), lr=learning_rate)
total_train_step = 0total_test_step = 0epoch = 30
writer = SummaryWriter("../logs_train")
for i in range(epoch):    print("--第 {} 輪訓練開始--".format(i+1))
        huaji.train()    for data in train_dataloader:        imgs, targets = data        imgs = imgs.to(device)        targets = targets.to(device)        outputs = huaji(imgs)        loss = loss_fn(outputs, targets)
                optimizer.zero_grad()        loss.backward()        optimizer.step()
        total_train_step = total_train_step + 1        if total_train_step % 100 == 0:            print("訓練次數：{}, Loss: {}".format(total_train_step, loss.item()))            writer.add_scalar("train_loss", loss.item(), total_train_step)
        huaji.eval()    total_test_loss = 0    total_accuracy = 0    with torch.no_grad():        for data in test_dataloader:            imgs, targets = data            imgs = imgs.to(device)            targets = targets.to(device)            outputs = huaji(imgs)            loss = loss_fn(outputs, targets)            total_test_loss = total_test_loss + loss.item()            accuracy = (outputs.argmax(1) == targets).sum()            total_accuracy = total_accuracy + accuracy
    print("整體測試集上的Loss: {}".format(total_test_loss))    print("整體測試集上的正確率: {}".format(total_accuracy/test_data_size))    writer.add_scalar("test_loss", total_test_loss, total_test_step)    writer.add_scalar("test_accuracy", total_accuracy/test_data_size, total_test_step)    total_test_step = total_test_step + 1
    torch.save(huaji, "huaji_{}.pth".format(i))    print("模型已保存")
writer.close()
使用GPU會大大提升模型的訓練速度，我們這裡就是用GPU進行訓練，如果你的PC上無法使用GPU,可以將代碼第10行的參數改為CPU。93行以上的代碼都是我們往期一一敲過的，我們把它組合起來。第97行代碼的作用是去掉數據的梯度，因為這裡我們是使用測試集給我們的模型打分，即給出當前的準確率。梯度在這裡不僅沒有作用，還會佔用我們的電腦資源。
我們把跑出的模型在測試機裡過一遍，得出數據，然後再與正確的結果進行對比，計算出判斷正確的個數之後，就得到了正確率。104行的item()，的佔用是把tensor類型轉化為int類型。在我們的驗證代碼裡還會用到它。
import torchimport torchvisionfrom PIL import Imagefrom torch import nn

c1 = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']image_path = r"E:\pycharm\code\learn_torch\imgs\04.png"image = Image.open(image_path)print(image)image = image.convert('RGB')transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),                                            torchvision.transforms.ToTensor()])
image = transform(image)print(image.shape)
class Huaji(nn.Module):    def __init__(self):        super(Huaji, self).__init__()        self.model = nn.Sequential(            nn.Conv2d(3, 32, 5, 1, 2),            nn.MaxPool2d(2),            nn.Conv2d(32, 32, 5, 1, 2),            nn.MaxPool2d(2),            nn.Conv2d(32, 64, 5, 1, 2),            nn.MaxPool2d(2),            nn.Flatten(),            nn.Linear(64*4*4, 64),            nn.Linear(64, 10)        )
    def forward(self, x):        x = self.model(x)        return x
model = torch.load("huaji_29.pth", map_location=torch.device('cpu'))print(model)image = torch.reshape(image, (1, 3, 32, 32))model.eval()with torch.no_grad():    output = model(image)print(output)
print(c1[output.argmax(1).item()])
注意我們需要把輸入圖片轉換成要求的格式，因此我們需要將圖片轉換成3通道且32*32大小的圖片，我們的模型要求輸入是4維的，我們通過reshape使其為1 batch_size。這樣就完成了對維度的要求，代碼第12行的作用是將圖片轉換成3通道，在學習opencv時，我們就了解到，有些圖片的格式是4通道，除了RGB外還有另一個Alpha Channel即透明度通道，我們需要去掉它。這裡的模型我們跑了30個循環，我們來看一下運行結果。
可以看到在測試集上的正確率是65%多一些，我們是十分類，這樣的正確率並不算差。我最開始只進行了10次循環，正確率在55%左右，可以看到增加訓練的次數，可以提高我們模型判斷的正確率。但要記住並不是增加足夠多的次數，正確率就能達到100%，它會有一個瓶頸，當它卡住時，只增加訓練次數並不會為我們的模型帶來更多的優化，這時我們就需要考慮其他方面，如增加訓練樣本，對模型自身進行優化，或是調整學習率等等。
接下來我們隨機在網上找一些圖片，這裡我找了6張如下。
分別是汽車、輪船、貓、飛機、飛機、貓。接著我們把上面的圖片交給我們訓練好的模型，依次得出以下結果。為了方便，我把tensor進行了相應的轉換。實際上返回的是tensor數據裡最大的值對應的位置。是訓練代碼第105行，前半部分代碼的作用
我們的模型判斷它們為汽車、輪船、鹿、鳥、飛機、貓。我們發現4，5張圖片判斷錯誤，分別把貓判斷成了鹿，把飛機判斷成了鳥。我們的正確率是65%多一些，你可以試著增加訓練次數，或是增加訓練樣本，看一下模型的正確率是否會提高。
到此為止，我們完成了一次完整的模型訓練，並進行了實際的測試。