tensorflow、多GPU、多线程训练VGG19来做cifar-10分类

背景：几天前需要写个多GPU训练的算法模型，翻来覆去在tensorflow的官网上看到cifar-10的官方代码，花了三天时间去掉代码的冗余部分和改写成自己的风格。

代码共有6部分构成：

1、data_input.py 由于cifar-10官方训练集和验证集都是.bin格式文件，故使用同一份数据读入代码

2、network.py 搭建VGG19，返回带weight decay的变量loss和交叉熵之和作为总loss

3、train.py 在每个GPU中建立tower，并行训练

4、val.py 多线程进行模型验证

5、toTFRecords.py 由于使用多线程无法读入label，故将图像和图像名（作label）制作成TFRecords

6、 test_tfrecords.py 读TFRecords文件，计算模型输出成csv

1、data_input.py：坑点为多线程读入数据时，如果num_threads==1，则万事大吉，否则，由于不同线程读取数据快慢不同，读入num_threads个数据的时候会出现一定程度的顺序打乱

import os
import tensorflow as tf

class data_input(object):
  def __init__(self, data_dir, batch_size, num_classes, is_training):
    self.data_dir = data_dir
    self.batch_size = batch_size
    self.num_classes = num_classes
    self.is_training = is_training
    
    self.image_batch, self.label_batch = self.load_data()
    
  #Input of this function is args, output is batch of images and batch of labels.
  def load_data(self):
    
    if self.is_training == True:
      filenames = [os.path.join(self.data_dir, 'data_batch_%d.bin' % i) for i in range(1, 6)]
    else:
      filenames = [os.path.join(self.data_dir, 'test_batch.bin')]
      
    for f in filenames:
      if not tf.gfile.Exists(f):
        raise ValueError('Failed to find file: ' + f)
    
    filename_queue = tf.train.string_input_producer(filenames, shuffle = False)
    image, label = self.load_one_sample(filename_queue)
    
    image.set_shape([32, 32, 3])
    label.set_shape([self.num_classes])
    
    #if the "num_threads" is not 1, due to the speed difference of each thread, there will be a shuffle in every num_threads data.
    if self.is_training == True:#Then data augmentation == True, shuffle == True.
      image = self.data_augmentation(image)
      image_batch, label_batch = tf.train.shuffle_batch([image, label], batch_size = self.batch_size, num_threads = 16, capacity = 20400, min_after_dequeue = 20000)
    else:
      image = tf.image.resize_image_with_crop_or_pad(image, 28, 28)
      image_batch, label_batch = tf.train.batch([image, label], batch_size = self.batch_size, num_threads = 16, capacity = 20400)
    
    return image_batch, tf.reshape(label_batch, [self.batch_size, self.num_classes])

  #From filename queue to read image and label. Image occupies 32*32*3 bytes, label occupies 1 bytes.
  def load_one_sample(self, filename_queue):
    
    image_bytes = 32 * 32 * 3
    label_bytes = 1
    
    record_bytes = image_bytes + label_bytes
    reader = tf.FixedLengthRecordReader(record_bytes = record_bytes)
    
    key, value = reader.read(filename_queue)
    record_bytes = tf.decode_raw(value, tf.uint8)
    
    label = tf.cast(tf.strided_slice(record_bytes, [0], [label_bytes]), tf.int32)
    label = tf.one_hot(label, self.num_classes)
    
    image = tf.reshape(tf.strided_slice(record_bytes, [label_bytes], [label_bytes + image_bytes]), [3, 32, 32])
    image = tf.transpose(image, [1, 2, 0])
    image = tf.cast(image, tf.float32)
    
    return image, tf.reshape(label, [self.num_classes])
    
  
  def data_augmentation(self, image):
    
    image = tf.random_crop(image, [28, 28, 3])
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_brightness(image, max_delta = 63)
    image = tf.image.random_contrast(image, lower = 0.2, upper = 1.8)
    image = tf.image.per_image_standardization(image)
    
    return image

2、network.py：坑点为不要忘记在验证集和测试集上关闭dropout。。。。

import tensorflow as tf

class NETWORK(object):
  def __init__(self, image_batch, label_batch, keep_prob, num_classes, is_training):
    self.image_batch = image_batch
    self.label_batch = label_batch
    
    if is_training is True:
      self.keep_prob = keep_prob
    else:
      self.keep_prob = 1
      
    self.num_classes = num_classes
    
    self.logits = self.inference()
    self.losses = self.loss()
  
  
  def inference(self):
    conv1_1 = conv(self.image_batch, 3, 64, 1, 'SAME', 0.0004, 'conv1_1')
    conv1_2 = conv(conv1_1, 3, 64, 2, 'SAME', None, 'conv1_2')
    pool1 = max_pool(conv1_2, 3, 2, 'SAME', 'pool1')
    
    conv2_1 = conv(pool1, 3, 128, 1, 'SAME', 0.0004, 'conv2_1')
    conv2_2 = conv(conv2_1, 3, 128, 1, 'SAME', None, 'conv2_2')
    pool2 = max_pool(conv2_2, 3, 2, 'SAME', 'pool2')

    conv3_1 = conv(pool2, 3, 256, 1, 'SAME', 0.0004, 'conv3_1')
    conv3_2 = conv(conv3_1, 3, 256, 1, 'SAME', None, 'conv3_2')
    conv3_3 = conv(conv3_2, 3, 256, 1, 'SAME', None, 'conv3_3')
    conv3_4 = conv(conv3_3, 3, 256, 1, 'SAME', None, 'conv3_4')
    pool3 = max_pool(conv3_4, 3, 2, 'SAME', 'pool3') 
    
    conv4_1 = conv(pool3, 3, 512, 1, 'SAME', 0.0004, 'conv4_1')
    conv4_2 = conv(conv4_1, 3, 512, 1, 'SAME', None, 'conv4_2')
    conv4_3 = conv(conv4_2, 3, 512, 1, 'SAME', None, 'conv4_3')
    conv4_4 = conv(conv4_3, 3, 512, 1, 'SAME', None, 'conv4_4')
    pool4 = max_pool(conv4_4, 3, 2, 'SAME', 'pool4')
    
    conv5_1 = conv(pool4, 3, 512, 1, 'SAME', 0.0004, 'conv5_1')
    conv5_2 = conv(conv5_1, 3, 512, 1, 'SAME', None, 'conv5_2')
    conv5_3 = conv(conv5_2, 3, 512, 1, 'SAME', None, 'conv5_3')
    conv5_4 = conv(conv5_3, 3, 512, 1, 'SAME', None, 'conv5_4')
    pool5 = max_pool(conv5_4, 3, 2, 'SAME', 'pool5') 
        
    reshape = tf.reshape(pool5, [pool5.get_shape().as_list()[0], -1])
    dim = reshape.get_shape()[1].value
        
    fc6 = fc(reshape, dim, 4096, 0.0004, 'fc6')
    dropout6 = dropout(fc6, self.keep_prob)
    
    fc7 = fc(dropout6, 4096, 4096, 0.0004, 'fc7')
    dropout7 = dropout(fc7, self.keep_prob)
    
    fc8 = fc(dropout7, 4096, self.num_classes, None, 'fc8')
    logits = tf.nn.softmax(fc8)
    
    return logits
    
  
  def loss(self):
    cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels = self.label_batch, logits = self.logits, name = 'cross_entropy_per_example')
    cross_entropy_mean = tf.reduce_mean(cross_entropy, name = 'cross_entropy')
    tf.add_to_collection('losses', cross_entropy_mean)
    return tf.add_n(tf.get_collection('losses'), name = 'total_loss')
    
    
def _variable_on_cpu(name, shape, initializer):
  with tf.device('/cpu:0'):
    var = tf.get_variable(name, shape, initializer = initializer, dtype = tf.float32)
  return var


def _variable_with_weight_decay(name, shape, stddev, wd):
  var = _variable_on_cpu(name, shape, tf.truncated_normal_initializer(stddev = stddev, dtype = tf.float32))
  if wd is not None:
    weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name = 'weight_loss')
    tf.add_to_collection('losses', weight_decay)
  return var


def conv(x, filter_wh, out_channels, stride_yx, padding, wd, name):
  in_channels = int(x.get_shape()[-1])
  
  with tf.variable_scope(name) as scope:
    kernel = _variable_with_weight_decay('weights', shape = [filter_wh, filter_wh, in_channels, out_channels], stddev = 5e-2, wd = wd)
    conv = tf.nn.conv2d(x, kernel, [1, stride_yx, stride_yx, 1], padding = padding)
    biases = _variable_on_cpu('biases', [out_channels], tf.constant_initializer(0.0))
    pre_activation = tf.nn.bias_add(conv, biases)
    conv_result = tf.nn.relu(pre_activation, name=scope.name)
    
  return conv_result


def max_pool(x, ksize_hw, stride_yx, padding, name):
  return tf.nn.max_pool(x, ksize = [1, ksize_hw, ksize_hw, 1], strides = [1, stride_yx, stride_yx, 1], padding = padding, name = name)


def lrn(x, depth_radius, alpha, beta, bias, name):
  return tf.nn.local_response_normalization(x, depth_radius, bias = bias, alpha = alpha, beta = beta, name = name)


def fc(x, num_in, num_out, wd, name):
  
  with tf.variable_scope(name) as scope:
    weights = _variable_with_weight_decay('weights', shape = [num_in, num_out], stddev = 0.04, wd = wd)
    biases = _variable_on_cpu('biases', [num_out], tf.constant_initializer(0.1))
    fc = tf.nn.xw_plus_b(x, weights, biases)
    fc_result = tf.nn.relu(fc, name = scope.name)
  
  return fc_result


def dropout(x, keep_prob):
  return tf.nn.dropout(x, keep_prob)

3、train.py：暂无坑点，多GPU并行训练就是在每个GPU上建立tower，并分别计算每个GPU中每个tower的loss和gradient，保留5个checkpoint。

import argparse
import sys, os
import re
import tensorflow as tf
from datetime import datetime
import time

from data_input import data_input
from network import NETWORK

#Calculate the total loss on a tower running the model.
def tower_loss(scope, images, labels, keep_prob, num_classes, is_training):
  model = NETWORK(images, labels, keep_prob, num_classes, is_training)
  _ = model.loss()
  
  #By using tf.add_to_collection() tf.get_collection() tf.add_n() to assemble all the losses.
  losses = tf.get_collection('losses', scope)
  total_loss = tf.add_n(losses, name='total_loss')
  
  for l in losses + [total_loss]:
    loss_name = re.sub('TOWER_[0-9]*/', '', l.op.name)
    tf.summary.scalar(loss_name, l)
    
  return total_loss

#Calculate the average gradient for each shared variable across all towers.
def average_gradients(tower_grads):
  average_grads = []
  
  for grad_and_vars in zip(*tower_grads):
    grads = []
    
    for g, _ in grad_and_vars:
      expanded_g = tf.expand_dims(g, 0)
      grads.append(expanded_g)

    grad = tf.concat(axis=0, values=grads)
    grad = tf.reduce_mean(grad, 0)
    v = grad_and_vars[0][1]
    grad_and_var = (grad, v)
    average_grads.append(grad_and_var)
    
  return average_grads


def train(args):
  #CPU is utilized to calculate "mean""update""variable statement".
  #GPU is utilized to calculate "loss""gradient".
  with tf.Graph().as_default(), tf.device('/cpu: 0'):
    
    #global_step == number of processed batches * number of GPU
    global_step = tf.get_variable('global_step', [], initializer = tf.constant_initializer(0), trainable = False)
    lr = tf.train.exponential_decay(args.base_lr, global_step, 10000, 0.9, staircase = True)
    opt = tf.train.GradientDescentOptimizer(lr)
    
    samples = data_input(args.train_data, args.batch_size, args.num_classes, args.is_training)
    batch_queue = tf.contrib.slim.prefetch_queue.prefetch_queue([samples.image_batch, samples.label_batch], capacity = 2 * args.num_gpus)
    
    tower_grads = []
    with tf.variable_scope(tf.get_variable_scope()):
      for i in range(args.num_gpus):
        with tf.device('/gpu: %d' %i):
          with tf.name_scope('TOWER_%d' % i) as scope:
            images, labels = batch_queue.dequeue()
            loss = tower_loss(scope, images, labels, args.dropout_rate, args.num_classes, args.is_training)
            tf.get_variable_scope().reuse_variables()
            summaries = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)
            grads = opt.compute_gradients(loss)
            tower_grads.append(grads)
            
    grads = average_gradients(tower_grads)
    summaries.append(tf.summary.scalar('learning_rate', lr))
    
    for grad, var in grads:
      if grad is not None:
        summaries.append(tf.summary.histogram(var.op.name + '/gradients', grad))
        
    apply_gradients_op = opt.apply_gradients(grads, global_step = global_step)
    
    for var in tf.trainable_variables():
      summaries.append(tf.summary.histogram(var.op.name, var))
    
    #Track the moving averages of all trainable variables.
    variable_averages = tf.train.ExponentialMovingAverage(0.9999, global_step)
    variable_averages_op = variable_averages.apply(tf.trainable_variables())
    
    #Group all updates to a single train op.
    train_op = tf.group(apply_gradients_op, variable_averages_op)
    
    #saver to save checkpoint file, summary_op to write summary file.
    saver = tf.train.Saver(tf.global_variables(), max_to_keep = 5)
    summary_op = tf.summary.merge(summaries)
    
    init = tf.global_variables_initializer()
    sess = tf.Session(config = tf.ConfigProto(allow_soft_placement = True, log_device_placement = False))
    sess.run(init)
    
    #Start queue runner to let data in.
    tf.train.start_queue_runners(sess = sess)
    summary_writer = tf.summary.FileWriter(args.train_logs, sess.graph)
    
    for step in range(args.max_iter):
      start_time = time.time()
      _, loss_value = sess.run([train_op, loss])
      duration = time.time() - start_time
      
      if step % 10 == 0:
        num_examples_per_step = args.batch_size * args.num_gpus
        examples_per_sec = num_examples_per_step / duration
        sec_per_batch = duration / args.num_gpus

        format_str = ('%s: step %d, loss = %.2f (%.1f examples/sec; %.3f sec/batch)')
        print (format_str % (datetime.now(), step, loss_value, examples_per_sec, sec_per_batch))

      if step % 100 == 0:
        summary_str = sess.run(summary_op)
        summary_writer.add_summary(summary_str, step)

      if step % 1000 == 0 or (step + 1) == args.max_iter:
        checkpoint_path = os.path.join(args.ckpt_dir, 'model.ckpt')
        saver.save(sess, checkpoint_path, global_step=step)


def main(args):
  #Clean the ancient summary file.
  if tf.gfile.Exists(args.train_logs):
    tf.gfile.DeleteRecursively(args.train_logs)
  tf.gfile.MakeDirs(args.train_logs)
  
  train(args)


def parse_arguments(argv):
  
  parser = argparse.ArgumentParser()
  
  parser.add_argument('--train_data', type = str, help = 'Train data directory.', default = './train_data')
  
  parser.add_argument('--ckpt_dir', type = str, help = 'Directory of trained models.', default = './ckpts')
  parser.add_argument('--train_logs', type = str, help = 'Directory for saving summary.', default = './train_logs')
  
  parser.add_argument('--num_gpus', type = int, help = 'Number of GPU.', default = 1)
  parser.add_argument('--max_iter', type = int, help = 'Iteration number for training.', default = 100000)
  parser.add_argument('--batch_size', type = int, help = 'Batch size.', default = 128)
  parser.add_argument('--num_classes', type = int, help = 'Number of classes.', default = 10)
  parser.add_argument('--base_lr', type = float, help = 'Base learning rate to start with.', default = 0.01)
  parser.add_argument('--weight_decay', type = float, help = 'Weight decay for L2 loss.', default = 0.0001)
  parser.add_argument('--momentum', type = float, help = 'Momentum for optimization.', default = 0.9)
  parser.add_argument('--dropout_rate', type = float, help = 'Dropout rate.', default = 0.7)
  
  parser.add_argument('--is_training', type = bool, help = 'Training mode.', default = True)
  
  return parser.parse_args(argv)


if __name__ == '__main__':
  main(parse_arguments(sys.argv[1: ]))

4、val.py：有个巨大的坑点。。。。那就是千万不要用tf.nn.in_top_k()函数，因为当预测值都为0时，该函数会显示你所有情况都预测对了。。。。

from datetime import datetime
import math, sys
import argparse

import numpy as np
import tensorflow as tf

from data_input import data_input
from network import NETWORK


def validation(args):
  
  with tf.Graph().as_default() as g:
    samples = data_input(args.val_data, args.batch_size, args.num_classes, args.is_training)
    model = NETWORK(samples.image_batch, samples.label_batch, args.dropout_rate, args.num_classes, args.is_training)
    logits = model.logits
    
    #Never use tf.nn.in_top_k() function!
    correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(samples.label_batch, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, 'float'))

    #To restore the moving averaged variabele.
    variable_averages = tf.train.ExponentialMovingAverage(0.9999)
    variables_to_restore = variable_averages.variables_to_restore()
    saver = tf.train.Saver(variables_to_restore)

    summary_op = tf.summary.merge_all()
    summary_writer = tf.summary.FileWriter(args.val_logs, g)

    with tf.Session() as sess:
      
      #coord is utilized to coordinate each thread.
      coord = tf.train.Coordinator()
      try:
        threads = []
        for qr in tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS):
          threads.extend(qr.create_threads(sess, coord = coord, daemon = True, start = True))
        
        #Try to remain top three model by running model on validation dataset.
        top_three_ckpt = [[1000, 0.1], [2000, 0.05], [3000, 0.0]]
        
        for global_step in range(2000, 7000, 1000):
          ckpt_path = './ckpts/model.ckpt-' + str(global_step)
          saver.restore(sess, ckpt_path)
    
          num_iter = int(math.ceil(10000 / args.batch_size))
          true_count = 0
          step = 0
          
          while step < num_iter and not coord.should_stop():
            predictions = sess.run([accuracy])
            true_count += np.sum(predictions)
            step += 1
    
          precision = true_count / num_iter
          print('%s: precision @ 1 = %.3f' % (datetime.now(), precision))
    
          for i in range(len(top_three_ckpt)):
            if precision >= top_three_ckpt[i][1]:
              top_three_ckpt.insert(i, [int(global_step), precision])
              top_three_ckpt.pop()
              break
            else:
              continue
          
          print(top_three_ckpt)

          summary = tf.Summary()
          summary.ParseFromString(sess.run(summary_op))
          summary.value.add(tag='Precision @ 1', simple_value=precision)
          summary_writer.add_summary(summary, global_step)
          
      except Exception as e:
        coord.request_stop(e)
    
      coord.request_stop()
      coord.join(threads, stop_grace_period_secs=10)


def main(args):
  if tf.gfile.Exists(args.val_logs):
    tf.gfile.DeleteRecursively(args.val_logs)
  tf.gfile.MakeDirs(args.val_logs)
  
  validation(args)


def parse_arguments(argv):
  
  parser = argparse.ArgumentParser()
  
  parser.add_argument('--val_data', type = str, help = 'Validation data directory.', default = './val_data')
  
  parser.add_argument('--ckpt_dir', type = str, help = 'Directory of trained models.', default = './ckpts')
  parser.add_argument('--val_logs', type = str, help = 'Directory for saving evaluation summary.', default = './val_logs')
  
  parser.add_argument('--num_gpus', type = int, help = 'Number of GPU.', default = 1)
  parser.add_argument('--max_iter', type = int, help = 'Iteration number for training.', default = 100000)
  parser.add_argument('--batch_size', type = int, help = 'Batch size.', default = 128)
  parser.add_argument('--num_classes', type = int, help = 'Number of classes.', default = 10)
  parser.add_argument('--base_lr', type = float, help = 'Base learning rate to start with.', default = 0.01)
  parser.add_argument('--weight_decay', type = float, help = 'Weight decay for L2 loss.', default = 0.0001)
  parser.add_argument('--momentum', type = float, help = 'Momentum for optimization.', default = 0.9)
  parser.add_argument('--dropout_rate', type = float, help = 'Dropout rate.', default = 0.7)
  
  parser.add_argument('--is_training', type = bool, help = 'Training mode.', default = False)
  
  return parser.parse_args(argv)


if __name__ == '__main__':
  main(parse_arguments(sys.argv[1: ]))

5、toTFRecords.py：暂无坑点，网上有关制作TFRecords格式文件一大堆。。。。就只想说，使用多线程tf.train.string_input_producer()函数是无法读取图像的label的！！！

import tensorflow as tf
from PIL import Image

classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

image_path = "./test_data"
TFRecord_path = "./test.tfrecords"
csv_path = "./trainLabels.csv"

def main():
  writer = tf.python_io.TFRecordWriter(TFRecord_path)
  
  for i in range(1, 121, 1):
    label_int = i
    image_name = image_path + '/' + str(i) + '.png'
    image = Image.open(image_name)
    image_raw = image.tobytes()
    
    example = tf.train.Example(features=tf.train.Features(feature={
        "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[label_int])),
        'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_raw]))
    }))
    
    writer.write(example.SerializeToString())
    
  writer.close()

if __name__ == '__main__':
  main()

6、test_tfrecords.py：暂无坑点，老老实实读入tfrecords得到预测logits并写入csv。

import sys, csv
import argparse
import tensorflow as tf

from network import NETWORK

dataset = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

#Read and decode the tfrecords format file.
def read_and_decode(filename):
    filename_queue = tf.train.string_input_producer([filename], shuffle = False)

    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    features = tf.parse_single_example(serialized_example, features={'label': tf.FixedLenFeature([], tf.int64), 'img_raw' : tf.FixedLenFeature([], tf.string),})

    img = tf.decode_raw(features['img_raw'], tf.uint8)
    img = tf.reshape(img, [32, 32, 3])
    img = tf.cast(img, tf.float32) * (1. / 255) - 0.5
    label = tf.cast(features['label'], tf.int32)
    
    return img, label


def test(args):
  
  with tf.Graph().as_default():
    #Read image batch and label batch.
    img, label = read_and_decode("/home/gz04025/Desktop/Syntax/test.tfrecords")
    img_batch, label_batch = tf.train.batch([img, label], batch_size = args.batch_size, capacity = 2000)
    
    #input the data to network and get the corresponding logits.
    init = tf.initialize_all_variables()
    model = NETWORK(img_batch, label_batch, args.dropout_rate, args.num_classes, args.is_training)
    logits = model.logits
    num_logits = tf.argmax(logits, 1)

    #Restore moving average variables.
    variable_averages = tf.train.ExponentialMovingAverage(0.9999)
    variables_to_restore = variable_averages.variables_to_restore()
    saver = tf.train.Saver(variables_to_restore)

    with tf.Session() as sess:
      sess.run(init)
      
      #Restore the saved model.
      ckpt = tf.train.get_checkpoint_state(args.ckpt_dir)
      if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)
        global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
      else:
        print('No checkpoint file found.')
        return
            
      coord = tf.train.Coordinator()
      try:
        threads = []
        for qr in tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS):
          threads.extend(qr.create_threads(sess, coord = coord, daemon = True, start = True))

        iter_num = 0
        while iter_num < (10000 / args.batch_size) and not coord.should_stop():
          numpy_num_logits = num_logits.eval(session = sess)
          str_logits = []
          for val in numpy_num_logits:
            str_logits.append(dataset[val])    
          
          iter_num += 1
                  
        with open('./test.csv', 'w', newline = '') as csvfile:
          writer = csv.writer(csvfile)
          writer.writerow(['id', 'label'])

          for i in range(len(str_logits)):
            writer.writerow([int(global_step) * args.batch_size + i + 1, str_logits[i]])
            
        csvfile.close()
        
      except Exception as e:
        coord.request_stop(e)
  
      coord.request_stop()
      coord.join(threads, stop_grace_period_secs=10)
      

def main(args):
  test(args)


def parse_arguments(argv):
  
  parser = argparse.ArgumentParser()
  
  parser.add_argument('--test_data', type = str, help = 'Test data directory.', default = './test_data')
  
  parser.add_argument('--ckpt_dir', type = str, help = 'Directory of trained models.', default = './ckpts')
  parser.add_argument('--test_result', type = str, help = 'Directory for saving test result.', default = './')
  
  parser.add_argument('--num_gpus', type = int, help = 'Number of GPU.', default = 1)
  parser.add_argument('--max_iter', type = int, help = 'Iteration number for training.', default = 100000)
  parser.add_argument('--batch_size', type = int, help = 'Batch size.', default = 12)
  parser.add_argument('--num_classes', type = int, help = 'Number of classes.', default = 10)
  parser.add_argument('--base_lr', type = float, help = 'Base learning rate to start with.', default = 0.01)
  parser.add_argument('--weight_decay', type = float, help = 'Weight decay for L2 loss.', default = 0.0001)
  parser.add_argument('--momentum', type = float, help = 'Momentum for optimization.', default = 0.9)
  parser.add_argument('--dropout_rate', type = float, help = 'Dropout rate.', default = 0.7)
  
  parser.add_argument('--is_training', type = bool, help = 'Training mode.', default = False)
  
  return parser.parse_args(argv)

if __name__ == '__main__':
  main(parse_arguments(sys.argv[1: ]))