PocketFlow自定义模型和数据集

- 日理万妓 2022-02-23 13:20 256阅读 0赞

### 文章目录 ###

*  自定义模型
 *  要点概括
 *  数据输入管道
 *  网络定义
 *  执行脚本
 *  PocketFlow中训练网络

# 自定义模型 #

整合自定义模型需要实现一个`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