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()