使用pytorch編寫貓狗識別軟件
數據下載
一、下載數據集並創建以下形式文件目錄
train.py: 用於創建並訓練模型，並生成訓練完成的參數文件。
setting.py: 用於存放訓練配置、超參數，包括學習率，訓練次數，裁剪圖片大小，每次訓練圖片數量，參數保存地址。
train: 存放下載的數據集（共25000張圖片，其中貓狗各12500張）。
func: 自定義包，存放部分操作。
├─setting.py
├─train.py
├─train
│ ├─cat
│ └─dog
└─func
└─__init__.py
二、拆分數據，划分訓練集和驗證集
在func目錄下新建get_address.py文件:
"""get_address""" def get_address(): """返回狗地址列表、貓地址列表、工作目錄""" import os print('{:-^30}'.format('數據集')) data_file = os.listdir('./train/') print('圖片或文件數量:', str(len(data_file))) # 25000 dog_file = list(filter(lambda x: x[:3] == 'dog', data_file)) cat_file = list(filter(lambda x: x[:3] == 'cat', data_file)) print('狗:', str(len(dog_file)), '\n貓:', str(len(cat_file))) # 狗:12500 貓:12500 root = os.getcwd() print('工作目錄:', root) # 工作目錄: L:\kaggle print('{:-^30}'.format('')) return dog_file, cat_file, root
獲取貓狗圖片的地址並返回。（其實，這里可與下面的arrange.py寫在一起，還不太熟練）。
在func目錄下新建arrange.py文件：
"""arrange.py""" def arrange(): """整理數據，移動圖片位置""" import shutil import os from .get_address import get_address dog_file, cat_file, root = get_address() print('開始數據整理') # 新建文件夾 for i in ['dog', 'cat']: for j in ['train', 'val']: try: os.makedirs(os.path.join(root,j,i)) except FileExistsError as e: pass # 移動10%(1250)的狗圖到驗證集 for i, file in enumerate(dog_file): ori_path = os.path.join(root, 'train', file) if i < 0.9*len(dog_file): des_path = os.path.join(root, 'train', 'dog') else: des_path = os.path.join(root, 'val', 'dog') shutil.move(ori_path, des_path) # 移動10%(1250)的貓圖到驗證集 for i, file in enumerate(cat_file): ori_path = os.path.join(root, 'train', file) if i < 0.9*len(cat_file): des_path = os.path.join(root, 'train', 'cat') else: des_path = os.path.join(root, 'val', 'cat') shutil.move(ori_path, des_path) print('數據整理完成')
通過此調用函數將train內10%的貓狗圖片提取出來，並新建val目錄進行存放。原來的train將作為訓練集（22500張），val
作為訓練集（2500張）。
將上述函數在__init__.py中進行調用（下文類似調用操作省略）
import func.arrange import func.get_address arrange = func.arrange.arrange get_address = func.get_address.get_address
三、獲取圖片數據並轉化
在func目錄下新建get_data.py文件：
"""get_data.py"""
def get_data(input_size, batch_size): """獲取文件數據並轉換""" from torchvision import transforms from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader # 串聯多個圖片變換的操作（訓練集） # transforms.RandomResizedCrop(input_size) 先隨機采集，然后對裁剪得到的圖像縮放為同一大小 # RandomHorizontalFlip() 以給定的概率隨機水平旋轉給定的PIL的圖像 # transforms.ToTensor() 將圖片轉換為Tensor,歸一化至[0,1] # transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) 歸一化處理(平均數，標准偏差) transform_train = transforms.Compose([ transforms.RandomResizedCrop(input_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # 獲取訓練集（通過上面的方面操作） train_set = ImageFolder('train', transform=transform_train) # 封裝訓練集 train_loader = DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True) # 串聯多個圖片變換的操作（驗證集） transform_val = transforms.Compose([ transforms.Resize([input_size, input_size]), # 注意 Resize 參數是 2 維，和 RandomResizedCrop 不同 transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # 獲取驗證集（通過上面的方面操作） val_set = ImageFolder('val', transform=transform_val) # 封裝驗證集 val_loader = DataLoader(dataset=val_set, batch_size=batch_size, shuffle=False) # 輸出 return transform_train, train_set, train_loader, transform_val, val_set, val_loader
此處使用pytorch 中的 ImageFolder 可以直接讀取圖片集數據（第一個參數決定文件夾地址），但是每個圖片大小各異，且需要轉化為可識別的數據。需要對讀取的圖片進行變換操作（即transform參數），除圖片縮放外，還需進行歸一化處理以減小數據復雜度和方便數據處理。通過transforms.Compose函數，可將這些圖片變化操作串聯，並通過ImageFolder的調用，快速獲取到所需要的數據。
四、將上述操作整合至train.py文件中
填寫將所需的參數、設置填寫在setting.py中
"""setting.py"""
"""調整設置""" input_size = 224 # 裁剪圖片大小 batch_size = 128 # 一次訓練所選取的樣本數(直接影響到GPU內存的使用情況) save_path = './weights.pt' # 訓練參數儲存地址 lr = 1e-3 # 學習率（后面用） n_epoch = 10 # 訓練次數（后面用）
填寫train.py
import torch from torchvision import models from torch import nn import func as f from setting import input_size, batch_size, save_path, lr, n_epoch f.arrange() # 整理數據，移動圖片位置(若已經整理完成可注釋) # 獲取文件數據並轉換成參數集 transform_train, train_set, train_loader, transform_val, val_set, val_loader = f.get_data(input_size, batch_size) print('映射關系:', train_set.class_to_idx) # {'cat': 0, 'dog': 1} print('訓練集長度:', len(train_set.imgs)) # 22500 print('訓練集規格:', train_set[1][0].size()) # torch.Size([3, 224, 224])
五、構建卷積神經網絡
使用Resnet18模型（殘差網絡介紹）
device = f.device() # 選擇訓練模式(GPU) print('訓練模式:', device, '模式') # 殘差網絡(18指定的是帶有權重的18層，包括卷積層和全連接層，不包括池化層和BN層) # pretrained=True 使用預訓練模型 # 使用resnet18模型 transfer_model = models.resnet18(pretrained=True) for param in transfer_model.parameters(): # 屏蔽預訓練模型的權重，只訓練最后一層的全連接的權重 param.requires_grad = False # 修改最后一層維數，即把原來的全連接層替換成輸出維數為2的全連接層 # 提取fc層中固定的參數 dim = transfer_model.fc.in_features # 設置網絡中的全連接層為2 transfer_model.fc = nn.Linear(dim, 2) # 構建神經網絡 net = transfer_model.to(device)
此處使用Resnet18模型為基礎進行訓練，使用預訓練模型加速訓練，使得模型收斂更快。另外，由於我們需要處理的是二元分類問題，全連接層輸出維數需為2（損失使用交叉熵損失函數，激活使用softmax）。
在func目錄下新建device.py文件：
"""device.py"""
def device(): """自動選擇訓練模式，盡可能使用GPU進行運算""" import torch if torch.cuda.is_available(): return torch.device('cuda:0') else: return torch.device('cpu')
選擇訓練模式，也可直接使用 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 進行替代。
六、創建訓練和驗證函數
在func目錄下新建train.py文件
"""train.py""" def train(net, optimizer, device, criterion, train_loader): """訓練""" net.train() batch_num = len(train_loader) running_loss = 0.0 for i, data in enumerate(train_loader, start=1): # 將輸入傳入GPU(CPU) inputs, labels = data inputs, labels = inputs.to(device), labels.to(device) # 參數梯度置零、向前、反向、優化 optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() # 計算誤差並顯示 running_loss += loss.item() if i % 10 == 0: print('batch:{}/{} loss:{:.3f}'.format(i, batch_num, running_loss / 20)) running_loss = 0.0
optimizer.zero_grad():梯度置零（因為梯度計算是累加的）。
outputs = net(inputs):向前傳播，求出預測值。
loss = criterion(outputs, labels):計算損失。
loss.backward():反向傳播，計算當前梯度。
optimizer.step() :根據梯度更新網絡參數。
在func目錄下新建validate.py文件：
"""validate.py"""
def validate(net, device, val_loader): """驗證函數""" import torch net.eval() # 測試，需關閉dropout correct = 0 total = 0 with torch.no_grad(): for data in val_loader: images, labels = data images, labels = images.to(device), labels.to(device) outputs = net(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print('測試圖像的網絡精度: %d %%' % (100 * correct / total))
with torch.no_grad():驗證無需進行計算，停止跟蹤歷史記錄和使用內存。
outputs = net(images):將圖片數據通過神經網絡，得到輸出值。
_, predicted = torch.max(outputs.data, 1):規格化返回預測值。
七、確定優化器、損失函數，進行訓練並保存
# 分類問題——交叉熵損失函數 criterion = nn.CrossEntropyLoss() # 優化器——隨機梯度下降 # 學習率lr=10^-3; optimizer = torch.optim.SGD(net.fc.parameters(), lr=lr) # optimizer = torch.optim.Adam(net.parameters(), lr=lr) for epoch in range(n_epoch): print('第{}次訓練'.format(epoch+1)) f.train(net, optimizer, device, criterion, train_loader) f.validate(net, device, val_loader) # 僅保存模型參數 torch.save(net.state_dict(), save_path)
保存的模型參數儲存在save_path地址中
單次訓練后准確度可達95%，經過十次訓練，准確度達到97%。
八、單圖片驗證，並進行可視化操作
在根目錄新建兩個文件tk.py和test.py：
"""test.py""" def test(): from PIL import Image import torch from torchvision import models from torch import nn from setting import input_size, save_path from torchvision import transforms # ------------------------ 加載數據 --------------------------- # # 定義預訓練變換 transform_val = transforms.Compose([ transforms.Resize([input_size, input_size]), # 注意 Resize 參數是 2 維，和 RandomResizedCrop 不同 transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) class_names = ['0', '180', '270', '90'] # 這個順序很重要，要和訓練時候的類名順序一致 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # ------------------------ 載入模型並且訓練 --------------------------- # transfer_model = models.resnet18(pretrained=True) for param in transfer_model.parameters(): param.requires_grad = False dim = transfer_model.fc.in_features transfer_model.fc = nn.Linear(dim, 2) # 構建神經網絡 net = transfer_model.to(device) net.load_state_dict(torch.load(save_path)) net.eval() image_PIL = Image.open(r'test/image') image_tensor = transform_val(image_PIL) # 以下語句等效於 image_tensor = torch.unsqueeze(image_tensor, 0) image_tensor.unsqueeze_(0) # 沒有這句話會報錯 image_tensor = image_tensor.to(device) out = net(image_tensor) # 得到預測結果，並且從大到小排序 _, indices = torch.sort(out, descending=True) # 返回每個預測值的百分數 percentage = torch.nn.functional.softmax(out, dim=1)[0] * 100 if percentage[0] > percentage[1]: out = '此圖片有{:>4.1f}%可能是只貓'.format(percentage[0]) else: out = '此圖片有{:>4.1f}%可能是只狗'.format(percentage[1]) return out
"""tk.py""" from tkinter import * from tkinter import filedialog from PIL import Image, ImageTk if __name__ == "__main__": root = Tk() root.title('請選擇圖片') frame = Frame(root, bd=2, relief=SUNKEN) frame.grid_rowconfigure(0, weight=1) frame.grid_columnconfigure(0, weight=1) canvas = Canvas(frame, bd=0) canvas.grid(row=0, column=0, sticky=N+S+E+W) frame.pack(fill=BOTH, expand=1) def printcoords(): import shutil, os File = filedialog.askopenfilename(parent=root, initialdir=os.getcwd(), title='選擇圖片.') img = Image.open(File) out = img.resize((336, 192), Image.ANTIALIAS) # resize image with high-quality filename = ImageTk.PhotoImage(out) canvas.image = filename canvas.create_image(0, 0, anchor='nw', image=filename) # print(File) dir = os.getcwd() + r'\test' shutil.rmtree(dir, True) os.makedirs(os.path.join(os.getcwd(), 'test')) shutil.copyfile(File, dir + r'\image') from test import test result = test() print(result) root.title(result) Button(root, text='選擇', command=printcoords).pack() root.mainloop()
這部分不是重點，就不細說了。運行tk.py結果如下：
