hewumars
4/4/2019 - 10:10 AM
PocketFlow自定义模型和数据集
[TOC]
自定义模型
整合自定义模型需要实现一个ModelHelper类。ModelHelper包含数据输入管道和网络前向传播和损失函数的定义。使用自定义的ModelHelper，网络可以无限制的使用FullPrecLearner训练，或者使用其他如ChannelPrunedLearner和UniformQuantTFLearner学习器进行通道裁剪和量化。
要点概括
要完成训练需要提供下面两个组件：
- 数据输入管道
- 网络定义
ModelHelper是抽象类AbstractModelHelper的子类，它涉及来提供上面这些定义。在PocketFlow，已经提供了几个ModelHelper类来描述数据集和模型结构的不同组合。要使用自定义模型，需要实现新的ModelHelper类。此外，还需要一个执行脚本来调用这个新定义的ModelHelper类。
数据输入管道
首先，需要告诉PocketFlow如何解析数据文件。在examples/fmnist_dataset.py定义一个名为FmnistDataset类创建返回训练和测试子集的迭代器。
from dataset.abstract_dataset import AbstractDataset
FLAGS = tf.app.flags.FLAGS
#将图像和标签读入内存
def load_mnist(image_file,label_file):
#...
return image,labels
def parse_fn(image,label,is_train):
"""
输入：
进行相关预处理，resize，crop，等
输出：image：图像张量
label：one-hot标签张量
"""
return image,label
class FMnistDataset(AbstractDataset):
'''数据集管道类'''
def __init__(self,is_train):
super(FMnistDataset,self).__init__(is_train)
#...
if is_train:
self.batch_size = FLAGS.batch_size
image_file = os.path.join(data_dir, 'train images file name')
label_file = os.path.join(data_dir, "train lables file name")
else:
#测试集，同上
self.images, self.labels = load_mnist(iamge_file,label_file)
self.parse_fn = lambda x,y : parse_fn(x,y,is_train)
def build(self,enbl_trn_val_split=False):
"""
构建tf.data.Dataset()的迭代器
"""
# create a tf.data.Dataset() object from NumPy arrays
dataset = tf.data.Dataset.from_tensor_slices((self.images, self.labels))
dataset = dataset.map(self.parse_fn, num_parallel_calls=FLAGS.nb_threads)
# create iterators for training & validation subsets separately
if self.is_train and enbl_trn_val_split:
iterator_val = self.__make_iterator(dataset.take(FLAGS.nb_smpls_val))
iterator_trn = self.__make_iterator(dataset.skip(FLAGS.nb_smpls_val))
return iterator_trn, iterator_val
return self.__make_iterator(dataset)
def __make_iterator(self, dataset):
"""从tf.data.Dataset创建迭代器
Args:
* dataset: tf.data.Dataset object
Returns:
* iterator: iterator for the dataset
"""
dataset = dataset.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=FLAGS.buffer_size))
dataset = dataset.batch(self.batch_size)
dataset = dataset.prefetch(FLAGS.prefetch_size)
iterator = dataset.make_one_shot_iterator()
return iterator
在创建FMnistDataset类的对象时，应提供名为is_train的额外参数来切换训练子集和测试子集。数据文件可以存储在本地计算机或HDFS集群上，目录路径在路径配置文件中指定，例如： data_dir_local_fmnist = / home / user_name / datasets / Fashion-MNIST 构造函数从* .gz文件加载图像和标签，每个文件都存储在NumPy数组中。然后使用构造函数从这两个NumPy数组创建TensorFlow的数据集迭代器。特别是，如果enbl_trn_val_split和is_train都为True，则原始训练子集将分为两部分，一部分用于模型训练，另一部分用于验证。
网络定义
实现新的ModelHelper类，它的全部实现放在./nets,命名convnet_at_fmnist.py、
import tensorflow as tf
from nets.abstract_model_helper import AbstractModelHelper
from datasets.fmnist_datset import FMnistDataset
from utils.lrn_rate_utils import setup_lrn_rate_piecewise_constant
from utils.multi_gup_wrapper import MultiGpuWrapper as mgw
FLAGS = tf.app.flags.FLAGS
#训练周期比率，学习率初始化值，批数量，向量，损失衰减
tf.app.flags.DEFINE_float('nb_epochs_rat', 1.0, '# of training epochs\' ratio')
tf.app.flags.DEFINE_float('lrn_rate_init', 1e-1, 'initial learning rate')
tf.app.flags.DEFINE_float('batch_size_norm', 128, 'normalization factor of batch size')
tf.app.flags.DEFINE_float('momentum', 0.9, 'momentum coefficient')
tf.app.flags.DEFINE_float('loss_w_dcy', 3e-4, 'weight decaying loss\'s coefficient')
def forward_fn(inputs,data_format):
"""
前向推理函数
参数：
inputs:网络前向传播的输入数据
data_format:数据格式（通道在前or通道在后）
返回：
inputs:网络前向传播的输出结果
"""
if data_format == 'channel_first':
inputs = tf.transpose(inputs,[0,3,1,2])#NCHW->NHWC
#conv1
inputs = tf.layers.conv2d(inputs,32,[5,5],padding='same',data_format=data_format,activation=tf.nn.relu,name='conv1')
intpus = tf.layers.max_pooling2d(inputs,[2,2],2,data_format=data_format,name='pool1')
#conv2
inputs = tf.layers.conv2d(inputs,64,[5,5],padding='same',data_format=data_format,activation=tf.nn.relu,name='conv2')
intpus = tf.layers.max_pooling2d(inputs,[2,2],2,data_format=data_format,name='pool2')
#fc3
inputs = tf.layers.flatten(inputs,name='flatten')
inputs = tf.layers.dense(input,1024,activation=tf.nn.relu,name='fc3')
#fc4
inputs = tf.layers.dense(input,FLAGS.nb_classes,name='fc4')
inputs = tf.nn.softmax(inputs,name='softmax')
return inputs
class ModelHelper(AbstractModelHelper):
"""模型帮助器：为Fashion-MNIST数据集创建ConvNet模型"""
#1
def __init__(self):
#继承初始化
super(ModelHelper,self).__init__()
#初始化数据集
self.dataset_train = FMnistDataset(is_train=True)
self.dataset_eval = FMnistDataset(is_train=False)
#2
def build_dataset_train(self,enbl_trn_val_split=False):
#通过数据增强构建用于训练的数据子集
return self.data_train.build(enbl_trn_val_split)
#3
def build_dataset_eval(self):
#不使用数据增强，构建评估子集
return self.data_eval.build()
#4
def forward_train(self,inputs,data_format='channel_last'):
#训练时，前向运算
return forward_fn(inputs,data_format)
#5 trainable_var不懂
def calc_loss(self,labels,outputs,trainable_vars):
#计算损失（和额外的评估指标）
loss = tf.losses.softmax_cross_entropy(labels,outputs)
loss += FLAGS.loss_w_dcy * tf.add_n([tf.nn.l2_loss(var) for var in trainable_vars])
accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmx(labels,axis=1),tf.argmax(outputs,axis=1)),tf.float32))
metrics = {'accuracy':accuracy}
return loss, metrics
#6
def setup_lrn_rate(self,global_step):
#设置学习率和迭代次数
nb_epochs = 160 #训练周期
idxs_epoch = [40,80,120] #周期索引
decay_rates = [1.0,0.1,0.01,0.001] #衰减率
batch_size = FLAGS.batch_size * (1 if not FLAGS.enbl_multi_gpu else mgw.size)#单卡或多卡batch_size
lrn_rate = setup_lrn_rate_piecewise_constant(global_step,batch_size,idxs_epoch,decay_rates)#得到学习率
nb_iters = int(FLAGS.nb_smpls_train*nb_epochs*FLAGS.nb_epochs_rat / batch_size)#迭代次数
return lrn_rate,nb_iters
@property
def model_name(self):
"""Model's name."""
return 'convnet'
@property
def dataset_name(self):
"""Dataset's name."""
return 'fmnist'
在build_dataset_train和build_dataset_eval函数中，我们采用先前引入的FMnistDataset类来定义数据输入管道。网络前向传递计算在forward_train和forward_eval函数中定义，它们分别对应于训练和评估graph。训练与评估graph略有不同，例如batchnorm相关的操作。 calc_loss函数计算损失函数的值和额外的评估指标，例如分类准确性。最后，setup_lrn_rate函数定义学习速率策略以及训练迭代次数。
执行脚本
除了自定义的ModelHelper类之外，我们还需要一个执行脚本将其传递给相应的模型压缩组件来启动训练过程。下面是完整的实现（这应该放在“./nets”目录下，命名为“convnet_at_fmnist_run.py”）：
import traceback
import tensorflow as tf
from nets.convnet_at_fmnist import ModelHelper #我们自定义的模型类
from learners.learner_utils import create_learner #创建学习器，用来选择相对应的压缩方法
FLAGS = tf.app.flags.FLAGS
#tensorboard日志目录，使能多卡，学习器类型，执行模式，调试
tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory')
tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training')
tf.app.flags.DEFINE_string('learner', 'full-prec', 'learner\'s name')
tf.app.flags.DEFINE_string('exec_mode', 'train', 'execution mode: train / eval')
tf.app.flags.DEFINE_boolean('debug', False, 'debugging information')
def main(unused_argv):
#程序入口
try:
if FLAGS.debug:
tf.logging.set_verbosity(tf.logging.DEBUG)
else:
tf.logging.set_verbosity(tf.logging.INFO)
sm_writer = tf.summary.FileWriter(FLAGS.log_dir)
# 打印FLAGS的值
tf.logging.info('FLAGS:')
for key,value in FLAGS.flag_values_dict().items():
tf.logging.info('{}: {}'.format(key, value))
# 构建模型帮助器和学习器
model_helper = ModelHelper()
learner = create_learner(sm_writer,model_helper)
# 执行学习器
if FLAGS.exec_mode == 'train':
learner.train()
elif FLAGS.exec_mode == 'eval':
learner.download_model()
learner.evaluate()
else:
raise ValueError('unrecognized execution mode: ' + FLAGS.exec_mode)
# 正常退出
return 0
except ValueError:
traceback.print_exc()
return 1
if __name__ == '__main__':
tf.app.run()
PocketFlow中训练网络
使用FullPrecLearner全精度学习器训练模型指令：
./scripts/run_local.sh nets/convnet_at_fmnist_run.py --learner full-prec
使用UniformQuantTFLearner均匀量化TF学习器训练模型指令：
./scripts/run_local.sh nets/convnet_at_fmnist_run.py --learner uniform-tf