Quantization Aware Training TF 코드

Quantization Aware Training 이라는 컨셉으로 구글이 몇년전에 벌써 내놓은 아이디이다 

주로 conv/deconv layer 에 asymmetric quantization 을 simulation 하는 방식으로 TF 에서  fake layer 를 추가해주는 API 를 제공해준다.

아래 처럼 모델을 define 하면서 optimizer 전에 fake layer 를 추가해준다 . 

하지만 단순히 floating 으로 학습된 checkpoint 를 save.restore API 로 그냥 불러 오면 erorr 를 뱉어 낼것이다 

 
class Net(object):
    def __init__(self, use_gpu_1 = True, QAT =True):
        g = tf.Graph()
 
        self.x1 = tf.placeholder(tf.float32, [None, H, W, 3], name='x1')  # image1
        self.x2 = tf.placeholder(tf.float32, [None, H, W, 3], name='x2')  # image2
        self.x3 = tf.placeholder(tf.float32, [None, H, W, 2], name='x3')  # label
        self.lr = tf.placeholder(tf.float32, [], name='lr')  # lr
 
        with tf.variable_scope('conv'):
            # pdb.set_trace()
            concat1 = tf.concat([self.x1, self.x2], 3)
 
            ### Original FlowNet
            conv1 = slim.conv2d(concat1, 64, [7, 7], 2, scope='conv1')
            conv2 = slim.conv2d(conv1, 128, [5, 5], 2, scope='conv2')
            conv3 = slim.conv2d(conv2, 256, [5, 5], 2, scope='conv3')
 
            ### FlowNet-SD
            # conv1 = slim.conv2d(concat1, 64, [3, 3], 2, scope='conv1')
            # conv2 = slim.conv2d(conv1, 128, [3, 3], 2, scope='conv2')
            # conv3 = slim.conv2d(conv2, 256, [3, 3], 2, scope='conv3')
 
            conv3_1 = slim.conv2d(conv3, 256, [3, 3], 1, scope='conv3_1')
            conv4 = slim.conv2d(conv3_1, 512, [3, 3], 2, scope='conv4')
            conv4_1 = slim.conv2d(conv4, 512, [3, 3], 1, scope='conv4_1')
            conv5 = slim.conv2d(conv4_1, 512, [3, 3], 2, scope='conv5')
            conv5_1 = slim.conv2d(conv5, 512, [3, 3], 1, scope='conv5_1')
            conv6 = slim.conv2d(conv5_1, 1024, [3, 3], 2, scope='conv6')
            conv6_1 = slim.conv2d(conv6, 1024, [3, 3], 1, scope='conv6_1')
            predict6 = slim.conv2d(conv6_1, 2, [3, 3], 1, activation_fn=None, scope='pred6')
 
 
 
        with tf.variable_scope('deconv'):
            # 12 * 16 flow
            deconv5 = slim.conv2d_transpose(conv6_1, 512, [4, 4], 2, scope='deconv5')
            deconvflow6 = slim.conv2d_transpose(predict6, 2, [4, 4], 2, 'SAME', scope='deconvflow6')
            concat5 = tf.concat([conv5_1, deconv5, deconvflow6], 3, name='concat5')
            predict5 = slim.conv2d(concat5, 2, [3, 3], 1, 'SAME', activation_fn=None, scope='predict5')
            # 24 * 32 flow
            deconv4 = slim.conv2d_transpose(concat5, 256, [4, 4], 2, 'SAME', scope='deconv4')
            deconvflow5 = slim.conv2d_transpose(predict5, 2, [4, 4], 2, 'SAME', scope='deconvflow5')
            concat4 = tf.concat([conv4_1, deconv4, deconvflow5], 3, name='concat4')
            predict4 = slim.conv2d(concat4, 2, [3, 3], 1, 'SAME', activation_fn=None, scope='predict4')
            # 48 * 64 flow
            deconv3 = slim.conv2d_transpose(concat4, 128, [4, 4], 2, 'SAME', scope='deconv3')
            deconvflow4 = slim.conv2d_transpose(predict4, 2, [4, 4], 2, 'SAME', scope='deconvflow4')
            concat3 = tf.concat([conv3_1, deconv3, deconvflow4], 3, name='concat3')
            predict3 = slim.conv2d(concat3, 2, [3, 3], 1, 'SAME', activation_fn=None, scope='predict3')
            # 96 * 128 flow
            deconv2 = slim.conv2d_transpose(concat3, 64, [4, 4], 2, 'SAME', scope='deconv2')
            deconvflow3 = slim.conv2d_transpose(predict3, 2, [4, 4], 2, 'SAME', scope='deconvflow3')
            concat2 = tf.concat([conv2, deconv2, deconvflow3], 3, name='concat2')
            predict2 = slim.conv2d(concat2, 2, [3, 3], 1, 'SAME', activation_fn=None, scope='predict2')
            # 192 * 256 flow
            deconv1 = slim.conv2d_transpose(concat2, 64, [4, 4], 2, 'SAME', scope='deconv1')
            deconvflow2 = slim.conv2d_transpose(predict2, 2, [4, 4], 2, 'SAME', scope='deconvflow2')
            concat1 = tf.concat([conv1, deconv1, deconvflow2], 3, name='concat1')
            predict1 = slim.conv2d(concat1, 2, [3, 3], 1, 'SAME', activation_fn=None, scope='predict1')
 
 
 
        self.tvars = tf.trainable_variables()  # turn on if you want to check the variables
        # self.variables_names = [v.name for v in self.tvars]
 
        with tf.variable_scope('loss'):
 
            weight = [1.0/2, 1.0/4, 1.0/8, 1.0/16, 1.0/32, 1.0/32]
            flow6 = tf.image.resize_images(self.x3, [6, 8])
            loss6 = weight[5] * self.mean_loss(flow6, predict6)
            flow5 = tf.image.resize_images(self.x3, [12, 16])
            loss5 = weight[4] * self.mean_loss(flow5, predict5)
            flow4 = tf.image.resize_images(self.x3, [24, 32])
            loss4 = weight[3] * self.mean_loss(flow4, predict4)
            flow3 = tf.image.resize_images(self.x3, [48, 64])
            loss3 = weight[2] * self.mean_loss(flow3, predict3)
            flow2 = tf.image.resize_images(self.x3, [96, 128])
            loss2 = weight[1] * self.mean_loss(flow2, predict2)
            flow1 = tf.image.resize_images(self.x3, [192, 256])
            loss1 = weight[0] * self.mean_loss(flow1, predict1)
            self.loss = tf.add_n([loss6, loss5, loss4, loss3, loss2, loss1])
            tf.summary.scalar('loss6', loss6)
            tf.summary.scalar('loss5', loss5)
            tf.summary.scalar('loss4', loss4)
            tf.summary.scalar('loss3', loss3)
            tf.summary.scalar('loss2', loss2)
            tf.summary.scalar('loss1', loss1)
            tf.summary.scalar('loss', self.loss)
            # pdb.set_trace()
            # self.merged = tf.merge_all_summaries()
 
            self.merged = tf.summary.merge_all()
 
 
        if (QAT == True):
 
            ##tf.train.import_meta_graph(meta_file)
            #g = tf.get_default_graph()
            tf.contrib.quantize.create_training_graph(quant_delay=0)
 
 
        optimizer = tf.train.AdamOptimizer(self.lr)
        self.train_op = slim.learning.create_train_op(self.loss, optimizer)
 
        # gpu configuration
        # self.tf_config = tf.ConfigProto()
        # self.tf_config = tf.ConfigProto(device_count = {'GPU': 0})
        self.tf_config = tf.ConfigProto(allow_soft_placement=True)
        self.tf_config = tf.ConfigProto(log_device_placement=False)
        self.tf_config.gpu_options.allow_growth = True
 
        # if use_gpu_1:
        #     self.tf_config.gpu_options.visible_device_list = '2,3'
 
        # pdb.set_trace()
 
        # self.init_all = tf.initialize_all_variables()
        self.init_all = tf.global_variables_initializer()
 
        self.train_tensor = slim.learning.create_train_op(
            self.loss,
            optimizer=optimizer)
 
        self.g = tf.Graph()
 
    def mean_loss(self, gt, predict):
        loss = tf.reduce_mean(tf.abs(gt-predict))
 
        return loss

그래서 tensorflow 에서 이렇게 trainable 한 param 과 구조를 가져 와서 fake layer 를 restore 하는 방식으로 

checkpoint 를  loading  해서 weight 값들의 min /max 를 찾아 가면서 int 8 의 quanitzation 을 시뮬레이션 하는 과정이다. 

        variables_to_restore = slim.get_variables_to_restore()
        reuse_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
 
        # Rename the other variables with the function name_in_checkpoint
        re_use = {name_in_checkpoint(var):var for var in reuse_vars}
        # Create a Saver to restore the checkpoint, given the variables
 
        saver = tf.train.Saver(re_use, max_to_keep=max_to_keep)
 
 
        sess.run(model.init_all)
        saver.restore(sess, dir_restore)