Pytorch【60天修炼计划】之第一阶段——入门2

原文链接： https://zhuanlan.zhihu.com/p/29024978

之前写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} 
'.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的方式来执行训练或者测试。