Pytorch【60天修煉計劃】之第一階段——入門2

之前寫Pytorch的 代碼時候，只用了一個main函數就搞定了，其實這樣可讀性是非常差的，所以我從知乎上找到了一篇知乎文章將如何合理組織代碼的文章，現在把它轉化成自己的東西。

項目文件組織架構

程序文件的組織架構爲：

程序文件的組織結構：
├── checkpoints/
├── data/
│ ├──_init_.py
│ ├── dataset.py
│ └── get_data.sh
├── models/
│ ├── _init_.py
│ ├── AlexNet.py
│ ├── BasicModule.py
│ └── ResNet34.py
└── utils/
│ ├── _init_.py
│ └── visualize.py
├── config.py
├── main.py
├── requirements.txt
├── README.md

其中：

• checkpoints/： 用於保存訓練好的模型，可使程序在異常退出後仍能重新載入模型，恢復訓練
• data/：數據相關操作，包括數據預處理、dataset實現等
• models/：模型定義，可以有多個模型，例如AlexNet和ResNet34，一個模型對應一個文件
• utils/：可能用到的工具函數，在本次實驗中主要是封裝了可視化工具
• config.py：配置文件，所有可配置的變量都集中在此，並提供默認值
• main.py：主文件，訓練和測試程序的入口，可通過不同的命令來指定不同的操作和參數
• requirements.txt：程序依賴的第三方庫
• README.md：提供程序的必要說明

轉自：https://www.zhihu.com/collection/341906824 ，侵刪。

1 _init_.py 是什麼？？

可以看到，幾乎每個文件夾下都有_init_.py，一個目錄如果包含了_init_.py 文件，那麼它就變成了一個包（package）。_init_.py可以爲空，也可以定義包的屬性和方法，但其必須存在，其它程序才能從這個目錄中導入相應的模塊或函數。例如在data/文件夾下有_init_.py，則在main.py 中就可以

from data.dataset import DogCat

而如果在data/_init_.py中寫入

from .dataset import DogCat

則在main.py中就可以直接寫爲：

from data import DogCat

或者

import data;
dataset = data.DogCat

相比於from data.dataset import DogCat更加便捷。

2 數據加載

數據的相關處理主要保存在data/dataset.py中。關於數據加載的相關操作，其基本原理就是使用Dataset進行數據集的封裝，再使用Dataloader實現數據並行加載。

一般提供的數據有訓練集和測試集，而在實際使用中，還需要驗證集，這就要從訓練集中劃分一些作爲驗證集。對於訓練集，我們希望做一些數據增強處理，如隨機裁剪、隨機翻轉、加噪聲等，而驗證集和測試集則不需要。下面看dataset.py的代碼：

#作者：陳雲
#鏈接：https://zhuanlan.zhihu.com/p/29024978
#來源：知乎
#著作權歸作者所有，轉載請聯繫作者獲得授權。

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)
               
       # shuffle imgs
       np.random.seed(100)
       imgs = np.random.permutation(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:
       
           # 數據轉換操作，測試驗證和訓練的數據轉換有所區別           
           normalize = T.Normalize(mean = [0.485, 0.456, 0.406], 
                                    std = [0.229, 0.224, 0.225])

           # 測試集和驗證集不用數據增強
           if self.test or not train: 
               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)

3 模型定義

模型的定義主要保存在models/目錄下，其中BasicModule是對nn.Module的簡易封裝，提供快速加載和保存模型的接口。

#作者：陳雲
#鏈接：https://zhuanlan.zhihu.com/p/29024978
#來源：知乎

class BasicModule(t.nn.Module):
   '''
   封裝了nn.Module，主要提供save和load兩個方法
   '''

   def __init__(self,opt=None):
       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 NewModule

即可。

其它關於模型定義的注意事項，在上一章中已詳細講解，這裏就不再贅述，總結起來就是：

儘量使用nn.Sequential（比如AlexNet）
將經常使用的結構封裝成子Module（比如GoogLeNet的Inception結構，ResNet的Residual Block結構）
將重複且有規律性的結構，用函數生成（比如VGG的多種變體，ResNet多種變體都是由多個重複卷積層組成）
感興趣的 讀者可以看看在models/resnet34.py如何用不到80行的代碼(包括空行和註釋)實現resnet34。當然這些模型在torchvision中有實現，而且還提供了預訓練的權重，讀者可以很方便的使用：

import torchvision as tv
resnet34 = tv.models.resnet34(pretrained=True)

4 工具函數

在項目中，我們可能會用到一些helper方法，這些方法可以統一放在utils/文件夾下，需要使用時再引入。在本例中主要是封裝了可視化工具visdom的一些操作，其代碼如下，在本次實驗中只會用到plot方法，用來統計損失信息。

#coding:utf8
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.iteritems():
           self.plot(k, v)

   def img_many(self, d):
       for k, v in d.iteritems():
           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=unicode(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)
       '''
       self.vis.images(img_.cpu().numpy(),
                      win=unicode(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):
       '''
       self.function 等價於self.vis.function
       自定義的plot,image,log,plot_many等除外
       '''
       return getattr(self.vis, name)

5 配置文件

在模型定義、數據處理和訓練等過程都有很多變量，這些變量應提供默認值，並統一放置在配置文件中，這樣在後期調試、修改代碼或遷移程序時會比較方便，在這裏我們將所有可配置項放在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等）
• 模型參數

這樣我們在程序中就可以這樣使用：

import models
from config import DefaultConfig

opt = DefaultConfig()
lr = opt.lr
model = getattr(models, opt.model)
dataset = DogCat(opt.train_data_root)

這些都只是默認參數，在這裏還提供了更新函數，根據字典更新配置參數。

def parse(self, kwargs):
        '''
        根據字典kwargs 更新 config參數
        '''
        # 更新配置參數
        for k, v in kwargs.iteritems():
            if not hasattr(self, k):
                # 警告還是報錯，取決於你個人的喜好
                warnings.warn("Warning: opt has not attribut %s" %k)
            setattr(self, k, v)
            
        # 打印配置信息	
        print('user config:')
        for k, v in self.__class__.__dict__.iteritems():
            if not k.startswith('__'):
                print(k, getattr(self, k))
  

這樣我們在實際使用時，並不需要每次都修改config.py，只需要通過命令行傳入所需參數，覆蓋默認配置即可。
例如：

opt = DefaultConfig()
new_config = {‘lr’:0.1,‘use_gpu’:False}
opt.parse(new_config)
opt.lr == 0.1

6 main.py

在主程序main.py中，主要包含四個函數，其中三個需要命令行執行，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()

根據fire的使用方法，可通過python main.py --args=xx的方式來執行訓練或者測試。