pytorch入門與實踐學習筆記：chapter6 貓狗大戰
阿新 • • 發佈：2019-01-05
1. 程式的組成架構
前面提到過，程式主要包含以下功能：
- 模型定義
- 資料載入
- 訓練模型
- 訓練過程視覺化
- 測試
2.關於__init__.py
可以看到，幾乎每個資料夾下都有`__init__.py`，一個目錄如果包含了`__init__.py` 檔案，那麼它就變成了一個包（package）。`__init__.py`可以為空，也可以定義包的屬性和方法，但其必須存在，其它程式才能從這個目錄中匯入相應的模組或函式。例如在`data/`資料夾下有`__init__.py`，則在`main.py` 中就可以`from data.dataset import DogCat`。而如果在`__init__.py`中寫入`from .dataset import DogCat`，則在main.py中就可以直接寫為：`from data import DogCat`，或者`import data; dataset = data.DogCat`，相比於`from data.dataset import DogCat`更加便捷。
3.資料載入
主要基於Dataset封裝資料集，再使用Dataloader實現資料並行載入。資料載入分成三個部分：訓練集、驗證集和測試集。在訓練集部分需要進行一些資料預處理，如隨機裁剪、隨即翻轉、加噪聲等。這裡不同階段載入資料的方式是採用加一些判斷來區分。
import os from PIL import Image from torch.utils import data import numpy as np from torchvision import transforms as T class DogCat(data.Dataset): def __init__(self, root, transforms=None, train=True, test=False): ''' 目標：獲取所有圖片地址，並根據訓練、驗證、測試劃分資料 ''' self.test = test imgs = [os.path.join(root, img) for img in os.listdir(root)] # test1: data/test1/8973.jpg # train: data/train/cat.10004.jpg if self.test: imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2].split('/')[-1])) else: imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2])) imgs_num = len(imgs) # 劃分訓練、驗證集，驗證:訓練 = 3:7 if self.test: self.imgs = imgs elif train: self.imgs = imgs[:int(0.7*imgs_num)] else : self.imgs = imgs[int(0.7*imgs_num):] if transforms is None: # 資料轉換操作，測試驗證和訓練的資料轉換有所區別 # 對影象進行歸一化，將畫素值壓縮至[-1,1] normalize = T.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]) # 測試集和驗證集 if self.test or not train: # 壓縮影象大小至224*224,並轉化成tensor的格式 self.transforms = T.Compose([ T.Scale(224), T.CenterCrop(224), T.ToTensor(), normalize ]) # 訓練集 else : self.transforms = T.Compose([ T.Scale(256), T.RandomSizedCrop(224), T.RandomHorizontalFlip(), T.ToTensor(), normalize ]) def __getitem__(self, index): ''' 返回一張圖片的資料 對於測試集，沒有label，返回圖片id，如1000.jpg返回1000 ''' img_path = self.imgs[index] if self.test: label = int(self.imgs[index].split('.')[-2].split('/')[-1]) else: label = 1 if 'dog' in img_path.split('/')[-1] else 0 data = Image.open(img_path) data = self.transforms(data) return data, label def __len__(self): ''' 返回資料集中所有圖片的個數 ''' return len(self.imgs)
注意：將檔案讀取等費時操作放在`__getitem__`函式中，利用多程序加速。避免一次性將所有圖片都讀進記憶體，不僅費時也會佔用較大記憶體，而且不易進行資料增強等操作。
在使用時，我們可通過dataloader載入資料。
train_dataset = DogCat(opt.train_data_root, train=True) trainloader = DataLoader(train_dataset, batch_size = opt.batch_size, shuffle = True, num_workers = opt.num_workers) for ii, (data, label) in enumerate(trainloader): train()
4.定義模型
模型的定義主要儲存在`models/`目錄下，其中`BasicModule`是對`nn.Module`的簡易封裝，提供快速載入和儲存模型的介面。
class BasicModule(t.nn.Module):
'''
封裝了nn.Module，主要提供save和load兩個方法
'''
def __init__(self):
super(BasicModule,self).__init__()
self.model_name = str(type(self)) # 模型的預設名字
def load(self, path):
'''
可載入指定路徑的模型
'''
self.load_state_dict(t.load(path))
def save(self, name=None):
'''
儲存模型，預設使用“模型名字+時間”作為檔名，
如AlexNet_0710_23:57:29.pth
'''
if name is None:
prefix = 'checkpoints/' + self.model_name + '_'
name = time.strftime(prefix + '%m%d_%H:%M:%S.pth')
t.save(self.state_dict(), name)
return name在實際使用中，直接呼叫`model.save()`及`model.load(opt.load_path)`即可。
其它自定義模型一般繼承`BasicModule`，然後實現自己的模型。其中`AlexNet.py`實現了AlexNet，`ResNet34`實現了ResNet34。在`models/__init__py`中，程式碼如下：
from .AlexNet import AlexNet
from .ResNet34 import ResNet34
from models import AlexNet
# 或
import models
model = models.AlexNet()
# 或
import models
model = getattr('models', 'AlexNet')()新增模型後只需要在`models/__init__.py`中加上`from .new_module import new_module`即可。
5.pytorch 視覺化 visdom
封裝視覺化工具vidom，以下程式碼只使用了plot方法。
import visdom
import time
import numpy as np
class Visualizer(object):
'''
封裝了visdom的基本操作，但是你仍然可以通過`self.vis.function`
或者`self.function`呼叫原生的visdom介面
比如
self.text('hello visdom')
self.histogram(t.randn(1000))
self.line(t.arange(0, 10),t.arange(1, 11))
'''
def __init__(self, env='default', **kwargs):
self.vis = visdom.Visdom(env=env, **kwargs)
# 畫的第幾個數，相當於橫座標
# 儲存（’loss',23） 即loss的第23個點
self.index = {}
self.log_text = ''
def reinit(self, env='default', **kwargs):
'''
修改visdom的配置
'''
self.vis = visdom.Visdom(env=env, **kwargs)
return self
def plot_many(self, d):
'''
一次plot多個
@params d: dict (name, value) i.e. ('loss', 0.11)
'''
for k, v in d.items():
self.plot(k, v)
def img_many(self, d):
for k, v in d.items():
self.img(k, v)
def plot(self, name, y, **kwargs):
'''
self.plot('loss', 1.00)
'''
x = self.index.get(name, 0)
self.vis.line(Y=np.array([y]), X=np.array([x]),
win=name,
opts=dict(title=name),
update=None if x == 0 else 'append',
**kwargs
)
self.index[name] = x + 1
def img(self, name, img_, **kwargs):
'''
self.img('input_img', t.Tensor(64, 64))
self.img('input_imgs', t.Tensor(3, 64, 64))
self.img('input_imgs', t.Tensor(100, 1, 64, 64))
self.img('input_imgs', t.Tensor(100, 3, 64, 64), nrows=10)
!!! don't ~~self.img('input_imgs', t.Tensor(100, 64, 64), nrows=10)~~ !!!
'''
self.vis.images(img_.cpu().numpy(),
win=name,
opts=dict(title=name),
**kwargs
)
def log(self, info, win='log_text'):
'''
self.log({'loss':1, 'lr':0.0001})
'''
self.log_text += ('[{time}] {info} <br>'.format(
time=time.strftime('%m%d_%H%M%S'),\
info=info))
self.vis.text(self.log_text, win)
def __getattr__(self, name):
'''
自定義的plot,image,log,plot_many等除外
self.function 等價於self.vis.function
'''
return getattr(self.vis, name)6.配置模型引數
在模型定義、資料處理和訓練等過程都有很多變數，這些變數應提供預設值，並統一放置在配置檔案中，這樣在後期除錯、修改程式碼或遷移程式時會比較方便，在這裡我們將所有可配置項放在`config.py`中。
class DefaultConfig(object):
env = 'default' # visdom 環境
model = 'AlexNet' # 使用的模型，名字必須與models/__init__.py中的名字一致
train_data_root = './data/train/' # 訓練集存放路徑
test_data_root = './data/test1' # 測試集存放路徑
load_model_path = 'checkpoints/model.pth' # 載入預訓練的模型的路徑，為None代表不載入
batch_size = 128 # batch size
use_gpu = True # use GPU or not
num_workers = 4 # how many workers for loading data
print_freq = 20 # print info every N batch
debug_file = '/tmp/debug' # if os.path.exists(debug_file): enter ipdb
result_file = 'result.csv'
max_epoch = 10
lr = 0.1 # initial learning rate
lr_decay = 0.95 # when val_loss increase, lr = lr*lr_decay
weight_decay = 1e-4 # 損失函式可配置的引數主要包括：
- 資料集引數（檔案路徑、batch_size等）
- 訓練引數（學習率、訓練epoch等）
- 模型引數
7.定義執行函式 main.py
main 函式包括四個函式，三個需要命令執行，main.py的程式碼如下：
def train(**kwargs):
'''
訓練
'''
pass
def val(model, dataloader):
'''
計算模型在驗證集上的準確率等資訊，用以輔助訓練
'''
pass
def test(**kwargs):
'''
測試（inference）
'''
pass
def help():
'''
列印幫助的資訊
'''
print('help')
if __name__=='__main__':
import fire
fire.Fire()