基于PyTorch實(shí)現(xiàn)一個(gè)簡(jiǎn)單的CNN圖像分類器

'''
包括了各種數(shù)據(jù)集的讀取處理，以及圖像相關(guān)處理方法
'''
from torch.utils.data import Dataset
import torch
import os
import cv2
from Config import mycfg
import random
import numpy as np


class ImageClassifyDataset(Dataset):
    def __init__(self, imagedir, labelfile, classify_num, train=True):
    	'''
    	這里進(jìn)行一些初始化操作。
    	'''
        self.imagedir = imagedir
        self.labelfile = labelfile
        self.classify_num = classify_num
        self.img_list = []
        # 讀取標(biāo)簽
        with open(self.labelfile, 'r') as fp:
            lines = fp.readlines()
            for line in lines:
                filepath = os.path.join(self.imagedir, line.split(";")[0].replace('\\', '/'))
                label = line.split(";")[1].strip('\n')
                self.img_list.append((filepath, label))
        if not train:
            self.img_list = random.sample(self.img_list, 50)

    def __len__(self):
        return len(self.img_list)
        
    def __getitem__(self, item):
	    '''
	    這個(gè)函數(shù)是關(guān)鍵，通過item(索引)來取數(shù)據(jù)集中的數(shù)據(jù)，
	    一般來說在這里才將圖像數(shù)據(jù)加載入內(nèi)存，之前存的是圖像的保存路徑
	    '''
        _int_label = int(self.img_list[item][1])	# label直接用0,1,2,3,4...表示不同類別
        label = torch.tensor(_int_label,dtype=torch.long)
        img = self.ProcessImgResize(self.img_list[item][0])
        return img, label

    def ProcessImgResize(self, filename):
    	'''
    	對(duì)圖像進(jìn)行一些預(yù)處理
    	'''
        _img = cv2.imread(filename)
        _img = cv2.resize(_img, (mycfg.IMG_WIDTH, mycfg.IMG_HEIGHT), interpolation=cv2.INTER_CUBIC)
        _img = _img.transpose((2, 0, 1))
        _img = _img / 255
        _img = torch.from_numpy(_img)
        _img = _img.to(torch.float32)
        return _img

from torch.utils.data import DataLoader
from MyDataset import ImageClassifyDataset

dataset = ImageClassifyDataset(imagedir, labelfile, 10)
dataloader = DataLoader(dataset, batch_size=5, shuffle=True,num_workers=5)
for index, data in enumerate(dataloader):
	print(index)	# batch索引
	print(data)		# 一個(gè)batch的{img,label}

'''
包括了各種模型、自定義的loss計(jì)算方法、optimizer
'''
import torch.nn as nn


class Simple_CNN(nn.Module):
    def __init__(self, class_num):
        super(Simple_CNN, self).__init__()
        self.class_num = class_num
        self.conv1 = nn.Sequential(
            nn.Conv2d(		# input: 3,400,600
                in_channels=3,
                out_channels=8,
                kernel_size=5,
                stride=1,
                padding=2
            ),
            nn.Conv2d(
                in_channels=8,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2
            ),
            nn.AvgPool2d(2),  # 16,400,600 --> 16,200,300
            nn.BatchNorm2d(16),
            nn.LeakyReLU(),
            nn.Conv2d(
                in_channels=16,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2
            ),
            nn.Conv2d(
                in_channels=16,
                out_channels=8,
                kernel_size=5,
                stride=1,
                padding=2
            ),
            nn.AvgPool2d(2),  # 8,200,300 --> 8,100,150
            nn.BatchNorm2d(8),
            nn.LeakyReLU(),
            nn.Conv2d(
                in_channels=8,
                out_channels=8,
                kernel_size=3,
                stride=1,
                padding=1
            ),
            nn.Conv2d(
                in_channels=8,
                out_channels=1,
                kernel_size=3,
                stride=1,
                padding=1
            ),
            nn.AvgPool2d(2),  # 1,100,150 --> 1,50,75
            nn.BatchNorm2d(1),
            nn.LeakyReLU()
        )
        self.line = nn.Sequential(
            nn.Linear(
                in_features=50 * 75,
                out_features=self.class_num
            ),
            nn.Softmax()
        )

    def forward(self, x):
        x = self.conv1(x)
        x = x.view(-1, 50 * 75)
        y = self.line(x)
        return y

self.cnnmodel = Simple_CNN(mycfg.CLASS_NUM)
self.criterion = nn.CrossEntropyLoss()	# 交叉熵，標(biāo)簽應(yīng)該是0,1,2,3...的形式而不是獨(dú)熱的
self.optimizer = optim.SGD(self.cnnmodel.parameters(), lr=mycfg.LEARNING_RATE, momentum=0.9)

for each_epoch in range(mycfg.MAX_EPOCH):
            running_loss = 0.0
            self.cnnmodel.train()
            for index, data in enumerate(self.dataloader):
                inputs, labels = data
                outputs = self.cnnmodel(inputs)
                loss = self.criterion(outputs, labels)

                self.optimizer.zero_grad()	# 清理上一次循環(huán)的梯度
                loss.backward()	# 反向傳播
                self.optimizer.step()	# 更新參數(shù)
                running_loss += loss.item()
                if index % 200 == 199:
                    print("[{}] loss: {:.4f}".format(each_epoch, running_loss/200))
                    running_loss = 0.0
            # 保存每一輪的模型
            model_name = 'classify-{}-{}.pth'.format(each_epoch,round(all_loss/all_index,3))
            torch.save(self.cnnmodel,model_name)	# 保存全部模型

acc = 0.0
count = 0
self.cnnmodel = torch.load('mymodel.pth')
self.cnnmodel.eval()
for index, data in enumerate(dataloader_eval):
	inputs, labels = data   # 5,3,400,600  5,10
	count += len(labels)
	outputs = cnnmodel(inputs)
	_,predict = torch.max(outputs, 1)
	acc += (labels == predict).sum().item()
print("[{}] accurancy: {:.4f}".format(each_epoch, acc / count))

torch.save(cnnmodel.state_dict(), "my_resnet.pth")
cnnmodel = Simple_CNN()
cnnmodel.load_state_dict(torch.load("my_resnet.pth"))