Permalink
Cannot retrieve contributors at this time
Name already in use
A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch?
splitnet-wrn/train_split.py
Go to fileThis commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
383 lines (325 sloc)
18.9 KB
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
import os | |
from datetime import datetime | |
import time | |
import tensorflow as tf | |
import numpy as np | |
import sys | |
import select | |
from IPython import embed | |
from StringIO import StringIO | |
import matplotlib.pyplot as plt | |
import cifar100 | |
import resnet_split as resnet | |
# Dataset Configuration | |
tf.app.flags.DEFINE_string('data_dir', './cifar100/train_val_split', """Path to the CIFAR-100 data.""") | |
tf.app.flags.DEFINE_integer('num_classes', 100, """Number of classes in the dataset.""") | |
tf.app.flags.DEFINE_integer('num_train_instance', 45000, """Number of training images.""") | |
tf.app.flags.DEFINE_integer('num_val_instance', 5000, """Number of val images.""") | |
# Network Configuration | |
tf.app.flags.DEFINE_integer('batch_size', 90, """Number of images to process in a batch.""") | |
tf.app.flags.DEFINE_integer('num_residual_units', 2, """Number of residual block per group. | |
Total number of conv layers will be 6n+4""") | |
tf.app.flags.DEFINE_integer('k', 8, """Network width multiplier""") | |
tf.app.flags.DEFINE_integer('ngroups1', 1, """Grouping number on logits""") | |
tf.app.flags.DEFINE_integer('ngroups2', 1, """Grouping number on unit_3_x""") | |
tf.app.flags.DEFINE_integer('ngroups3', 1, """Grouping number on unit_2_x""") | |
# Optimization Configuration | |
tf.app.flags.DEFINE_float('l2_weight', 0.0001, """L2 loss weight applied all the weights""") | |
tf.app.flags.DEFINE_float('momentum', 0.9, """The momentum of MomentumOptimizer""") | |
tf.app.flags.DEFINE_float('initial_lr', 0.1, """Initial learning rate""") | |
tf.app.flags.DEFINE_string('lr_step_epoch', "80.0,120.0,160.0", """Epochs after which learing rate decays""") | |
tf.app.flags.DEFINE_float('lr_decay', 0.1, """Learning rate decay factor""") | |
tf.app.flags.DEFINE_boolean('finetune', False, """Whether to finetune.""") | |
# Training Configuration | |
tf.app.flags.DEFINE_string('train_dir', './train', """Directory where to write log and checkpoint.""") | |
tf.app.flags.DEFINE_integer('max_steps', 100000, """Number of batches to run.""") | |
tf.app.flags.DEFINE_integer('display', 100, """Number of iterations to display training info.""") | |
tf.app.flags.DEFINE_integer('val_interval', 1000, """Number of iterations to run a val""") | |
tf.app.flags.DEFINE_integer('val_iter', 100, """Number of iterations during a val""") | |
tf.app.flags.DEFINE_integer('checkpoint_interval', 10000, """Number of iterations to save parameters as a checkpoint""") | |
tf.app.flags.DEFINE_float('gpu_fraction', 0.95, """The fraction of GPU memory to be allocated""") | |
tf.app.flags.DEFINE_boolean('log_device_placement', False, """Whether to log device placement.""") | |
tf.app.flags.DEFINE_string('basemodel', None, """Base model to load paramters""") | |
tf.app.flags.DEFINE_string('checkpoint', None, """Model checkpoint to load""") | |
FLAGS = tf.app.flags.FLAGS | |
def get_lr(initial_lr, lr_decay, lr_decay_steps, global_step): | |
lr = initial_lr | |
for s in lr_decay_steps: | |
if global_step >= s: | |
lr *= lr_decay | |
return lr | |
def train(): | |
print('[Dataset Configuration]') | |
print('\tCIFAR-100 dir: %s' % FLAGS.data_dir) | |
print('\tNumber of classes: %d' % FLAGS.num_classes) | |
print('\tNumber of training images: %d' % FLAGS.num_train_instance) | |
print('\tNumber of val images: %d' % FLAGS.num_val_instance) | |
print('[Network Configuration]') | |
print('\tBatch size: %d' % FLAGS.batch_size) | |
print('\tResidual blocks per group: %d' % FLAGS.num_residual_units) | |
print('\tNetwork width multiplier: %d' % FLAGS.k) | |
print('\tNumber of Groups: %d-%d-%d' % (FLAGS.ngroups3, FLAGS.ngroups2, FLAGS.ngroups1)) | |
print('\tBasemodel file: %s' % FLAGS.basemodel) | |
print('[Optimization Configuration]') | |
print('\tL2 loss weight: %f' % FLAGS.l2_weight) | |
print('\tThe momentum optimizer: %f' % FLAGS.momentum) | |
print('\tInitial learning rate: %f' % FLAGS.initial_lr) | |
print('\tEpochs per lr step: %s' % FLAGS.lr_step_epoch) | |
print('\tLearning rate decay: %f' % FLAGS.lr_decay) | |
print('\tFinetune: %d' % FLAGS.finetune) | |
print('[Training Configuration]') | |
print('\tTrain dir: %s' % FLAGS.train_dir) | |
print('\tTraining max steps: %d' % FLAGS.max_steps) | |
print('\tSteps per displaying info: %d' % FLAGS.display) | |
print('\tSteps per validation: %d' % FLAGS.val_interval) | |
print('\tSteps during validation: %d' % FLAGS.val_iter) | |
print('\tSteps per saving checkpoints: %d' % FLAGS.checkpoint_interval) | |
print('\tGPU memory fraction: %f' % FLAGS.gpu_fraction) | |
print('\tLog device placement: %d' % FLAGS.log_device_placement) | |
with tf.Graph().as_default(): | |
init_step = 0 | |
global_step = tf.Variable(0, trainable=False, name='global_step') | |
# Get images and labels of CIFAR-100 | |
print('Load CIFAR-100 dataset') | |
train_dataset_path = os.path.join(FLAGS.data_dir, 'train') | |
val_dataset_path = os.path.join(FLAGS.data_dir, 'val') | |
print('\tLoading training data from %s' % train_dataset_path) | |
with tf.variable_scope('train_image'): | |
cifar100_train = cifar100.CIFAR100Runner(train_dataset_path, image_per_thread=32, | |
shuffle=True, distort=True, capacity=10000) | |
train_images, train_labels = cifar100_train.get_inputs(FLAGS.batch_size) | |
print('\tLoading validation data from %s' % val_dataset_path) | |
with tf.variable_scope('val_image'): | |
cifar100_val = cifar100.CIFAR100Runner(val_dataset_path, image_per_thread=32, | |
shuffle=False, distort=False, capacity=5000) | |
# shuffle=False, distort=False, capacity=10000) | |
val_images, val_labels = cifar100_val.get_inputs(FLAGS.batch_size) | |
# Build a Graph that computes the predictions from the inference model. | |
images = tf.placeholder(tf.float32, [FLAGS.batch_size, cifar100.IMAGE_SIZE, cifar100.IMAGE_SIZE, 3]) | |
labels = tf.placeholder(tf.int32, [FLAGS.batch_size]) | |
# Get splitted params | |
if not FLAGS.basemodel: | |
print('No basemodel found to load split params') | |
sys.exit(-1) | |
else: | |
print('Load split params from %s' % FLAGS.basemodel) | |
def get_perms(q_name, ngroups): | |
split_alpha = reader.get_tensor(q_name+'/alpha') | |
q_amax = np.argmax(split_alpha, axis=0) | |
return [np.where(q_amax == i)[0] for i in range(ngroups)] | |
reader = tf.train.NewCheckpointReader(FLAGS.basemodel) | |
split_params = {} | |
print('\tlogits...') | |
base_logits_w = reader.get_tensor('logits/fc/weights') | |
base_logits_b = reader.get_tensor('logits/fc/biases') | |
split_p1_idxs = get_perms('group/split_p1', FLAGS.ngroups1) | |
split_q1_idxs = get_perms('group/split_q1', FLAGS.ngroups1) | |
logits_params = {'weights':[], 'biases':[], 'input_perms':[], 'output_perms':[]} | |
for i in range(FLAGS.ngroups1): | |
logits_params['weights'].append(base_logits_w[split_p1_idxs[i], :][:, split_q1_idxs[i]]) | |
logits_params['biases'].append(base_logits_b[split_q1_idxs[i]]) | |
logits_params['input_perms'] = split_p1_idxs | |
logits_params['output_perms'] = split_q1_idxs | |
split_params['logits'] = logits_params | |
if FLAGS.ngroups2 > 1: | |
print('\tunit_3_x...') | |
base_unit_3_0_shortcut_k = reader.get_tensor('unit_3_0/shortcut/kernel') | |
base_unit_3_0_conv1_k = reader.get_tensor('unit_3_0/conv_1/kernel') | |
base_unit_3_0_conv2_k = reader.get_tensor('unit_3_0/conv_2/kernel') | |
base_unit_3_1_conv1_k = reader.get_tensor('unit_3_1/conv_1/kernel') | |
base_unit_3_1_conv2_k = reader.get_tensor('unit_3_1/conv_2/kernel') | |
split_p2_idxs = get_perms('group/split_p2', FLAGS.ngroups2) | |
split_q2_idxs = _merge_split_idxs(split_p1_idxs, _get_even_merge_idxs(FLAGS.ngroups1, FLAGS.ngroups2)) | |
split_r21_idxs = get_perms('group/split_r21', FLAGS.ngroups2) | |
split_r22_idxs = get_perms('group/split_r22', FLAGS.ngroups2) | |
unit_3_0_params = {'shortcut':[], 'conv1':[], 'conv2':[], 'p_perms':[], 'q_perms':[], 'r_perms':[]} | |
for i in range(FLAGS.ngroups2): | |
unit_3_0_params['shortcut'].append(base_unit_3_0_shortcut_k[:,:,split_p2_idxs[i],:][:,:,:,split_q2_idxs[i]]) | |
unit_3_0_params['conv1'].append(base_unit_3_0_conv1_k[:,:,split_p2_idxs[i],:][:,:,:,split_r21_idxs[i]]) | |
unit_3_0_params['conv2'].append(base_unit_3_0_conv2_k[:,:,split_r21_idxs[i],:][:,:,:,split_q2_idxs[i]]) | |
unit_3_0_params['p_perms'] = split_p2_idxs | |
unit_3_0_params['q_perms'] = split_q2_idxs | |
unit_3_0_params['r_perms'] = split_r21_idxs | |
split_params['unit_3_0'] = unit_3_0_params | |
unit_3_1_params = {'conv1':[], 'conv2':[], 'p_perms':[], 'r_perms':[]} | |
for i in range(FLAGS.ngroups2): | |
unit_3_1_params['conv1'].append(base_unit_3_1_conv1_k[:,:,split_q2_idxs[i],:][:,:,:,split_r22_idxs[i]]) | |
unit_3_1_params['conv2'].append(base_unit_3_1_conv2_k[:,:,split_r22_idxs[i],:][:,:,:,split_q2_idxs[i]]) | |
unit_3_1_params['p_perms'] = split_q2_idxs | |
unit_3_1_params['r_perms'] = split_r22_idxs | |
split_params['unit_3_1'] = unit_3_1_params | |
if FLAGS.ngroups3 > 1: | |
print('\tconv4_x...') | |
base_unit_2_0_shortcut_k = reader.get_tensor('unit_2_0/shortcut/kernel') | |
base_unit_2_0_conv1_k = reader.get_tensor('unit_2_0/conv_1/kernel') | |
base_unit_2_0_conv2_k = reader.get_tensor('unit_2_0/conv_2/kernel') | |
base_unit_2_1_conv1_k = reader.get_tensor('unit_2_1/conv_1/kernel') | |
base_unit_2_1_conv2_k = reader.get_tensor('unit_2_1/conv_2/kernel') | |
split_p3_idxs = get_perms('group/split_p3', FLAGS.ngroups3) | |
split_q3_idxs = _merge_split_idxs(split_p2_idxs, _get_even_merge_idxs(FLAGS.ngroups2, FLAGS.ngroups3)) | |
split_r31_idxs = get_perms('group/split_r31', FLAGS.ngroups3) | |
split_r32_idxs = get_perms('group/split_r32', FLAGS.ngroups3) | |
unit_2_0_params = {'shortcut':[], 'conv1':[], 'conv2':[], 'p_perms':[], 'q_perms':[], 'r_perms':[]} | |
for i in range(FLAGS.ngroups3): | |
unit_2_0_params['shortcut'].append(base_unit_2_0_shortcut_k[:,:,split_p3_idxs[i],:][:,:,:,split_q3_idxs[i]]) | |
unit_2_0_params['conv1'].append(base_unit_2_0_conv1_k[:,:,split_p3_idxs[i],:][:,:,:,split_r31_idxs[i]]) | |
unit_2_0_params['conv2'].append(base_unit_2_0_conv2_k[:,:,split_r31_idxs[i],:][:,:,:,split_q3_idxs[i]]) | |
unit_2_0_params['p_perms'] = split_p3_idxs | |
unit_2_0_params['q_perms'] = split_q3_idxs | |
unit_2_0_params['r_perms'] = split_r31_idxs | |
split_params['unit_2_0'] = unit_2_0_params | |
unit_2_1_params = {'conv1':[], 'conv2':[], 'p_perms':[], 'r_perms':[]} | |
for i in range(FLAGS.ngroups3): | |
unit_2_1_params['conv1'].append(base_unit_2_1_conv1_k[:,:,split_q3_idxs[i],:][:,:,:,split_r32_idxs[i]]) | |
unit_2_1_params['conv2'].append(base_unit_2_1_conv2_k[:,:,split_r32_idxs[i],:][:,:,:,split_q3_idxs[i]]) | |
unit_2_1_params['p_perms'] = split_q3_idxs | |
unit_2_1_params['r_perms'] = split_r32_idxs | |
split_params['unit_2_1'] = unit_2_1_params | |
# Build model | |
lr_decay_steps = map(float,FLAGS.lr_step_epoch.split(',')) | |
lr_decay_steps = map(int,[s*FLAGS.num_train_instance/FLAGS.batch_size for s in lr_decay_steps]) | |
hp = resnet.HParams(batch_size=FLAGS.batch_size, | |
num_classes=FLAGS.num_classes, | |
num_residual_units=FLAGS.num_residual_units, | |
k=FLAGS.k, | |
weight_decay=FLAGS.l2_weight, | |
ngroups1=FLAGS.ngroups1, | |
ngroups2=FLAGS.ngroups2, | |
ngroups3=FLAGS.ngroups3, | |
split_params=split_params, | |
momentum=FLAGS.momentum, | |
finetune=FLAGS.finetune) | |
network = resnet.ResNet(hp, images, labels, global_step) | |
network.build_model() | |
network.build_train_op() | |
print('Number of Weights: %d' % network._weights) | |
print('FLOPs: %d' % network._flops) | |
train_summary_op = tf.summary.merge_all() # Summaries(training) | |
# Build an initialization operation to run below. | |
init = tf.global_variables_initializer() | |
# Start running operations on the Graph. | |
sess = tf.Session(config=tf.ConfigProto( | |
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=FLAGS.gpu_fraction), | |
# allow_soft_placement=True, | |
log_device_placement=FLAGS.log_device_placement)) | |
'''debugging attempt | |
from tensorflow.python import debug as tf_debug | |
sess = tf_debug.LocalCLIDebugWrapperSession(sess) | |
def _get_data(datum, tensor): | |
return tensor == train_images | |
sess.add_tensor_filter("get_data", _get_data) | |
''' | |
sess.run(init) | |
# Create a saver. | |
saver = tf.train.Saver(tf.global_variables(), max_to_keep=10000) | |
if FLAGS.checkpoint is not None: | |
saver.restore(sess, FLAGS.checkpoint) | |
init_step = global_step.eval(session=sess) | |
print('Load checkpoint %s' % FLAGS.checkpoint) | |
else: | |
# Define a different saver to load model checkpoints | |
# Select only base variables (exclude split layers) | |
print('Load parameters from basemodel %s' % FLAGS.basemodel) | |
variables = tf.global_variables() | |
vars_restore = [var for var in variables | |
if not "Momentum" in var.name and | |
not "logits" in var.name and | |
not "global_step" in var.name] | |
if FLAGS.ngroups2 > 1: | |
vars_restore = [var for var in vars_restore | |
if not "unit_3_" in var.name] | |
if FLAGS.ngroups3 > 1: | |
vars_restore = [var for var in vars_restore | |
if not "unit_2_" in var.name] | |
saver_restore = tf.train.Saver(vars_restore, max_to_keep=10000) | |
saver_restore.restore(sess, FLAGS.basemodel) | |
# Start queue runners & summary_writer | |
cifar100_train.start_threads(sess, n_threads=20) | |
cifar100_val.start_threads(sess, n_threads=1) | |
if not os.path.exists(FLAGS.train_dir): | |
os.mkdir(FLAGS.train_dir) | |
summary_writer = tf.summary.FileWriter(os.path.join(FLAGS.train_dir, str(global_step.eval(session=sess))), | |
sess.graph) | |
# Training! | |
val_best_acc = 0.0 | |
for step in xrange(init_step, FLAGS.max_steps): | |
# val | |
if step % FLAGS.val_interval == 0: | |
val_loss, val_acc = 0.0, 0.0 | |
for i in range(FLAGS.val_iter): | |
val_images_val, val_labels_val = sess.run([val_images, val_labels]) | |
loss_value, acc_value = sess.run([network.loss, network.acc], | |
feed_dict={network.is_train:False, images:val_images_val, labels:val_labels_val}) | |
val_loss += loss_value | |
val_acc += acc_value | |
val_loss /= FLAGS.val_iter | |
val_acc /= FLAGS.val_iter | |
val_best_acc = max(val_best_acc, val_acc) | |
format_str = ('%s: (val) step %d, loss=%.4f, acc=%.4f') | |
print (format_str % (datetime.now(), step, val_loss, val_acc)) | |
val_summary = tf.Summary() | |
val_summary.value.add(tag='val/loss', simple_value=val_loss) | |
val_summary.value.add(tag='val/acc', simple_value=val_acc) | |
val_summary.value.add(tag='val/best_acc', simple_value=val_best_acc) | |
summary_writer.add_summary(val_summary, step) | |
summary_writer.flush() | |
# Train | |
lr_value = get_lr(FLAGS.initial_lr, FLAGS.lr_decay, lr_decay_steps, step) | |
start_time = time.time() | |
train_images_val, train_labels_val = sess.run([train_images, train_labels]) | |
_, loss_value, acc_value, train_summary_str = \ | |
sess.run([network.train_op, network.loss, network.acc, train_summary_op], | |
feed_dict={network.is_train:True, network.lr:lr_value, images:train_images_val, labels:train_labels_val}) | |
duration = time.time() - start_time | |
assert not np.isnan(loss_value) | |
# Display & Summary(training) | |
if step % FLAGS.display == 0: | |
num_examples_per_step = FLAGS.batch_size | |
examples_per_sec = num_examples_per_step / duration | |
sec_per_batch = float(duration) | |
format_str = ('%s: (Training) step %d, loss=%.4f, acc=%.4f, lr=%f (%.1f examples/sec; %.3f ' | |
'sec/batch)') | |
print (format_str % (datetime.now(), step, loss_value, acc_value, lr_value, | |
examples_per_sec, sec_per_batch)) | |
summary_writer.add_summary(train_summary_str, step) | |
# Save the model checkpoint periodically. | |
if (step > init_step and step % FLAGS.checkpoint_interval == 0) or (step + 1) == FLAGS.max_steps: | |
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt') | |
saver.save(sess, checkpoint_path, global_step=step) | |
if sys.stdin in select.select([sys.stdin], [], [], 0)[0]: | |
char = sys.stdin.read(1) | |
if char == 'b': | |
embed() | |
def _merge_split_q(q, merge_idxs, name='merge'): | |
ngroups, dim = q.shape | |
max_idx = np.max(merge_idxs) | |
temp_list = [] | |
for i in range(max_idx + 1): | |
temp = [] | |
for j in range(ngroups): | |
if merge_idxs[j] == i: | |
temp.append(q[j,:]) | |
temp_list.append(np.sum(temp, axis=0)) | |
ret = np.array(temp_list) | |
return ret | |
def _merge_split_idxs(split_idxs, merge_idxs, name='merge'): | |
ngroups = len(split_idxs) | |
max_idx = np.max(merge_idxs) | |
ret = [] | |
for i in range(max_idx + 1): | |
temp = [] | |
for j in range(ngroups): | |
if merge_idxs[j] == i: | |
temp.append(split_idxs[j]) | |
ret.append(np.concatenate(temp)) | |
return ret | |
def _get_even_merge_idxs(N, split): | |
assert N >= split | |
num_elems = [(N + split - i - 1)/split for i in range(split)] | |
expand_split = [[i] * n for i, n in enumerate(num_elems)] | |
return [t for l in expand_split for t in l] | |
def main(argv=None): # pylint: disable=unused-argument | |
train() | |
if __name__ == '__main__': | |
tf.app.run() |