""" Code for the MAML algorithm and network definitions. """
from __future__ import print_function
import numpy as np
import sys
import tensorflow as tf
try:
import special_grads
except KeyError as e:
print('WARN: Cannot define MaxPoolGrad, likely already defined for this version of tensorflow: %s' % e,
file=sys.stderr)
from collections import OrderedDict
from tensorflow.python.platform import flags
from utils import xent, conv_block, normalize, bn_relu_conv_block
FLAGS = flags.FLAGS
class MAML:
def __init__(self, dim_input=1, dim_output_train=1, dim_output_val=1, test_num_updates=5):
""" must call construct_model() after initializing MAML! """
self.dim_input = dim_input
self.dim_output_train = dim_output_train
self.dim_output_val = dim_output_val
self.update_lr = FLAGS.update_lr
self.meta_lr = tf.placeholder_with_default(FLAGS.meta_lr, ())
self.classification = False
self.test_num_updates = test_num_updates
self.loss_func = xent
self.classification = True
if FLAGS.on_encodings:
print('Meta-learning on encodings')
self.dim_hidden = [FLAGS.num_filters] * FLAGS.num_hidden_layers
print('hidden layers: {}'.format(self.dim_hidden))
self.forward = self.forward_fc
self.construct_weights = self.construct_fc_weights
else:
if FLAGS.conv:
self.dim_hidden = FLAGS.num_filters
if FLAGS.resnet:
if FLAGS.input_type == 'images_84x84':
self.forward = self.forward_resnet84
self.construct_weights = self.construct_resnet_weights84
assert FLAGS.num_parts_per_res_block == 2
assert FLAGS.num_res_blocks == 4
self.num_parts_per_res_block = FLAGS.num_parts_per_res_block
self.blocks = ['input', 'maxpool', 'res0', 'maxpool', 'res1', 'maxpool', 'res2', 'maxpool', 'res3', 'output']
elif FLAGS.input_type == 'images_224x224':
self.forward = self.forward_resnet224
self.construct_weights = self.construct_resnet_weights224
assert FLAGS.num_parts_per_res_block == 2
assert FLAGS.num_res_blocks == 4
self.num_parts_per_res_block = FLAGS.num_parts_per_res_block
self.blocks = ['input', 'maxpool', 'res0', 'maxpool', 'res1', 'maxpool', 'res2', 'maxpool', 'res3', 'output']
else:
raise ValueError
else:
self.forward = self.forward_conv
self.construct_weights = self.construct_conv_weights
else:
self.dim_hidden = [1024, 512, 256, 128]
print('hidden layers: {}'.format(self.dim_hidden))
self.forward = self.forward_fc
self.construct_weights = self.construct_fc_weights
if FLAGS.dataset == 'mnist' or FLAGS.dataset == 'omniglot':
self.channels = 1
else:
self.channels = 3
self.img_size = int(np.sqrt(self.dim_input/self.channels))
if FLAGS.dataset not in ['mnist', 'omniglot', 'miniimagenet', 'celeba', 'imagenet']:
raise ValueError('Unrecognized data source.')
# resnet things
def construct_model(self, input_tensors=None, prefix='metatrain_'):
# a: training data for inner gradient, b: test data for meta gradient
if prefix == 'metatrain_':
inner_update_batch_size = FLAGS.inner_update_batch_size_train
else:
inner_update_batch_size = FLAGS.inner_update_batch_size_val
outer_update_batch_size = FLAGS.outer_update_batch_size
if input_tensors is None:
self.inputa = tf.placeholder(tf.float32)
self.inputb = tf.placeholder(tf.float32)
self.labela = tf.placeholder(tf.float32)
self.labelb = tf.placeholder(tf.float32)
else:
self.inputa = input_tensors['inputa']
self.inputb = input_tensors['inputb']
self.labela = input_tensors['labela']
self.labelb = input_tensors['labelb']
if prefix == 'metaval_':
self.mv_inputa = self.inputa
self.mv_inputb = self.inputb
self.mv_labela = self.labela
self.mv_labelb = self.labelb
with tf.variable_scope('model', reuse=None) as training_scope:
if 'weights' in dir(self):
training_scope.reuse_variables()
weights = self.weights
else:
# Define the weights
self.weights = weights = self.construct_weights()
print(weights.keys())
# outputbs[i] and lossesb[i] is the output and loss after i+1 gradient updates
lossesa, outputas, lossesb, outputbs = [], [], [], []
accuraciesa, accuraciesb = [], []
num_updates = max(self.test_num_updates, FLAGS.num_updates)
outputbs = [[]]*num_updates
lossesb = [[]]*num_updates
accuraciesb = [[]]*num_updates
if FLAGS.from_scratch:
train_accuracies = [[]]*num_updates
def task_metalearn(inp, reuse=True):
""" Perform gradient descent for one task in the meta-batch. """
inputa, inputb, labela, labelb = inp
task_outputbs, task_lossesb = [], []
if FLAGS.from_scratch:
task_outputas = []
if self.classification:
task_accuraciesb = []
if FLAGS.from_scratch:
task_accuraciesa = []
task_outputa = self.forward(inputa, weights, prefix, reuse=reuse) # only reuse on the first iter
if FLAGS.from_scratch:
task_outputas.append(task_outputa)
task_lossa = self.loss_func(task_outputa, labela, inner_update_batch_size)
grads = tf.gradients(task_lossa, list(weights.values()))
if FLAGS.stop_grad:
grads = [tf.stop_gradient(grad) for grad in grads]
gradients = dict(zip(weights.keys(), grads))
fast_weights = dict(zip(weights.keys(), [weights[key] - self.update_lr*gradients[key] for key in weights.keys()]))
output = self.forward(inputb, fast_weights, prefix, reuse=True)
task_outputbs.append(output)
task_lossesb.append(self.loss_func(output, labelb, outer_update_batch_size))
for j in range(num_updates - 1):
outputa = self.forward(inputa, fast_weights, prefix, reuse=True)
loss = self.loss_func(outputa, labela, inner_update_batch_size)
if FLAGS.from_scratch:
task_outputas.append(outputa)
grads = tf.gradients(loss, list(fast_weights.values()))
if FLAGS.stop_grad:
grads = [tf.stop_gradient(grad) for grad in grads]
gradients = dict(zip(fast_weights.keys(), grads))
fast_weights = dict(zip(fast_weights.keys(), [fast_weights[key] - self.update_lr*gradients[key] for key in fast_weights.keys()]))
output = self.forward(inputb, fast_weights, prefix, reuse=True)
task_outputbs.append(output)
task_lossesb.append(self.loss_func(output, labelb, outer_update_batch_size))
task_output = [task_outputa, task_outputbs, task_lossa, task_lossesb]
if self.classification:
task_accuracya = tf.contrib.metrics.accuracy(tf.argmax(tf.nn.softmax(task_outputa), 1), tf.argmax(labela, 1))
for j in range(num_updates):
task_accuraciesb.append(tf.contrib.metrics.accuracy(tf.argmax(tf.nn.softmax(task_outputbs[j]), 1), tf.argmax(labelb, 1)))
if FLAGS.from_scratch:
task_accuraciesa.append(tf.contrib.metrics.accuracy(tf.argmax(tf.nn.softmax(task_outputas[j]), 1), tf.argmax(labela, 1)))
task_output.extend([task_accuracya, task_accuraciesb])
if FLAGS.from_scratch:
task_output.extend([task_accuraciesa])
return task_output
if FLAGS.norm is not 'None':
# to initialize the batch norm vars, might want to combine this, and not run idx 0 twice.
unused = task_metalearn((self.inputa[0], self.inputb[0], self.labela[0], self.labelb[0]), False)
out_dtype = [tf.float32, [tf.float32]*num_updates, tf.float32, [tf.float32]*num_updates]
if self.classification:
out_dtype.extend([tf.float32, [tf.float32]*num_updates])
if FLAGS.from_scratch:
out_dtype.extend([[tf.float32] * num_updates])
result = tf.map_fn(task_metalearn, elems=(self.inputa, self.inputb, self.labela, self.labelb), dtype=out_dtype, parallel_iterations=FLAGS.meta_batch_size)
if self.classification:
if FLAGS.from_scratch:
outputas, outputbs, lossesa, lossesb, accuraciesa, accuraciesb, train_accuracies = result
else:
outputas, outputbs, lossesa, lossesb, accuraciesa, accuraciesb = result
else:
outputas, outputbs, lossesa, lossesb = result
## Performance & Optimization
if 'train' in prefix:
self.total_loss1 = total_loss1 = tf.reduce_sum(lossesa) / tf.to_float(FLAGS.meta_batch_size)
self.total_losses2 = total_losses2 = [tf.reduce_sum(lossesb[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]
# after the map_fn
self.outputas, self.outputbs = outputas, outputbs
if self.classification:
self.total_accuracy1 = total_accuracy1 = tf.reduce_sum(accuraciesa) / tf.to_float(FLAGS.meta_batch_size)
self.total_accuracies2 = total_accuracies2 = [tf.reduce_sum(accuraciesb[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]
self.pretrain_op = tf.train.AdamOptimizer(self.meta_lr).minimize(total_loss1)
if FLAGS.metatrain_iterations > 0:
optimizer = tf.train.AdamOptimizer(self.meta_lr)
self.gvs = gvs = optimizer.compute_gradients(self.total_losses2[FLAGS.num_updates-1])
if FLAGS.dataset == 'miniimagenet' or FLAGS.dataset == 'celeba' or FLAGS.dataset == 'imagenet':
gvs = [(tf.clip_by_value(grad, -10, 10), var) for grad, var in gvs]
self.metatrain_op = optimizer.apply_gradients(gvs)
else:
self.metaval_total_loss1 = total_loss1 = tf.reduce_sum(lossesa) / tf.to_float(FLAGS.meta_batch_size)
self.metaval_total_losses2 = total_losses2 = [tf.reduce_sum(lossesb[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]
if self.classification:
self.metaval_total_accuracy1 = total_accuracy1 = tf.reduce_sum(accuraciesa) / tf.to_float(FLAGS.meta_batch_size)
self.metaval_total_accuracies2 = total_accuracies2 =[tf.reduce_sum(accuraciesb[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]
self.mv_outputbs = outputbs
if FLAGS.from_scratch:
self.metaval_train_accuracies = [tf.reduce_sum(train_accuracies[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]
## Summaries
tf.summary.scalar(prefix+'Pre-update loss', total_loss1)
if self.classification:
tf.summary.scalar(prefix+'Pre-update accuracy', total_accuracy1)
for j in range(num_updates):
tf.summary.scalar(prefix+'Post-update loss, step ' + str(j+1), total_losses2[j])
if self.classification:
tf.summary.scalar(prefix+'Post-update accuracy, step ' + str(j+1), total_accuracies2[j])
### Network construction functions (fc networks and conv networks)
def construct_fc_weights(self):
weights = {}
weights['w1'] = tf.Variable(tf.truncated_normal([self.dim_input, self.dim_hidden[0]], stddev=0.01))
weights['b1'] = tf.Variable(tf.zeros([self.dim_hidden[0]]))
for i in range(1,len(self.dim_hidden)):
weights['w'+str(i+1)] = tf.Variable(tf.truncated_normal([self.dim_hidden[i-1], self.dim_hidden[i]], stddev=0.01))
weights['b'+str(i+1)] = tf.Variable(tf.zeros([self.dim_hidden[i]]))
weights['w'+str(len(self.dim_hidden)+1)] = tf.Variable(tf.truncated_normal([self.dim_hidden[-1], self.dim_output_train], stddev=0.01))
weights['b'+str(len(self.dim_hidden)+1)] = tf.Variable(tf.zeros([self.dim_output_train]))
return weights
def forward_fc(self, inp, weights, prefix, reuse=False):
hidden = normalize(tf.matmul(inp, weights['w1']) + weights['b1'], activation=tf.nn.relu, reuse=reuse, scope='0')
for i in range(1,len(self.dim_hidden)):
hidden = normalize(tf.matmul(hidden, weights['w'+str(i+1)]) + weights['b'+str(i+1)], activation=tf.nn.relu, reuse=reuse, scope=str(i+1))
logits = tf.matmul(hidden, weights['w'+str(len(self.dim_hidden)+1)]) + weights['b'+str(len(self.dim_hidden)+1)]
if 'val' in prefix:
logits = tf.gather(logits, tf.range(self.dim_output_val), axis=1)
return logits
def construct_conv_weights(self):
weights = {}
dtype = tf.float32
conv_initializer = tf.contrib.layers.xavier_initializer_conv2d(dtype=dtype)
fc_initializer = tf.contrib.layers.xavier_initializer(dtype=dtype)
k = 3
channels = self.channels
weights['conv1'] = tf.get_variable('conv1', [k, k, channels, self.dim_hidden], initializer=conv_initializer, dtype=dtype)
weights['b1'] = tf.Variable(tf.zeros([self.dim_hidden]))
weights['conv2'] = tf.get_variable('conv2', [k, k, self.dim_hidden, self.dim_hidden], initializer=conv_initializer, dtype=dtype)
weights['b2'] = tf.Variable(tf.zeros([self.dim_hidden]))
weights['conv3'] = tf.get_variable('conv3', [k, k, self.dim_hidden, self.dim_hidden], initializer=conv_initializer, dtype=dtype)
weights['b3'] = tf.Variable(tf.zeros([self.dim_hidden]))
weights['conv4'] = tf.get_variable('conv4', [k, k, self.dim_hidden, self.dim_hidden], initializer=conv_initializer, dtype=dtype)
weights['b4'] = tf.Variable(tf.zeros([self.dim_hidden]))
if FLAGS.dataset == 'miniimagenet' or FLAGS.dataset == 'celeba' or FLAGS.dataset == 'imagenet':
# assumes max pooling
weights['w5'] = tf.get_variable('w5', [self.dim_hidden*5*5, self.dim_output_train], initializer=fc_initializer)
weights['b5'] = tf.Variable(tf.zeros([self.dim_output_train]), name='b5')
else:
weights['w5'] = tf.Variable(tf.random_normal([self.dim_hidden, self.dim_output_train]), name='w5')
weights['b5'] = tf.Variable(tf.zeros([self.dim_output_train]), name='b5')
return weights
def forward_conv(self, inp, weights, prefix, reuse=False, scope=''):
# reuse is for the normalization parameters.
channels = self.channels
inp = tf.reshape(inp, [-1, self.img_size, self.img_size, channels])
hidden1 = conv_block(inp, weights['conv1'], weights['b1'], reuse, scope+'0')
hidden2 = conv_block(hidden1, weights['conv2'], weights['b2'], reuse, scope+'1')
hidden3 = conv_block(hidden2, weights['conv3'], weights['b3'], reuse, scope+'2')
hidden4 = conv_block(hidden3, weights['conv4'], weights['b4'], reuse, scope+'3')
if FLAGS.dataset == 'miniimagenet' or FLAGS.dataset == 'celeba' or FLAGS.dataset == 'imagenet':
# last hidden layer is 6x6x64-ish, reshape to a vector
hidden4 = tf.reshape(hidden4, [-1, np.prod([int(dim) for dim in hidden4.get_shape()[1:]])])
else:
hidden4 = tf.reduce_mean(hidden4, [1, 2])
logits = tf.matmul(hidden4, weights['w5']) + weights['b5']
if 'val' in prefix:
logits = tf.gather(logits, tf.range(self.dim_output_val), axis=1)
return logits
def construct_resnet_weights224(self):
weights = OrderedDict()
dtype = tf.float32
conv_initializer = tf.contrib.layers.xavier_initializer_conv2d(dtype=dtype)
bias_initializer = tf.zeros_initializer(dtype=dtype)
fc_initializer = tf.contrib.layers.xavier_initializer(dtype=dtype)
def make_conv_layer_weights(weights, scope, k, filters_in, filters_out, bias=True):
weights['{}/conv'.format(scope)] = tf.get_variable('{}/conv'.format(scope), [k, k, filters_in, filters_out], initializer=conv_initializer, dtype=dtype)
if bias:
weights['{}/bias'.format(scope)] = tf.get_variable('{}/bias'.format(scope), [filters_out], initializer=bias_initializer, dtype=dtype)
def make_fc_layer_weights(weights, scope, dims_in, dims_out):
weights['{}/fc'.format(scope)] = tf.get_variable('{}/fc'.format(scope), [dims_in, dims_out], initializer=fc_initializer, dtype=dtype)
weights['{}/bias'.format(scope)] = tf.get_variable('{}/bias'.format(scope), [dims_out], initializer=bias_initializer, dtype=dtype)
for block_name in self.blocks:
if block_name == 'input':
make_conv_layer_weights(weights, block_name, k=7, filters_in=self.channels, filters_out=64)
elif 'res' in block_name:
j = int(block_name[-1])
last_block_filter = 64 if j == 0 else 64 * 2 ** (j-1)
this_block_filter = 64 if j == 0 else last_block_filter * 2
print(block_name, last_block_filter, this_block_filter)
make_conv_layer_weights(weights, '{}/shortcut'.format(block_name), k=1, filters_in=last_block_filter,
filters_out=this_block_filter, bias=False)
for i in range(self.num_parts_per_res_block):
make_conv_layer_weights(weights, '{}/part{}'.format(block_name, i), k=3,
filters_in=last_block_filter if i == 0 else this_block_filter,
filters_out=this_block_filter)
elif block_name == 'output':
make_fc_layer_weights(weights, block_name, dims_in=512, dims_out=self.dim_output_train)
return weights
def construct_resnet_weights84(self):
weights = OrderedDict()
dtype = tf.float32
conv_initializer = tf.contrib.layers.xavier_initializer_conv2d(dtype=dtype)
bias_initializer = tf.zeros_initializer(dtype=dtype)
fc_initializer = tf.contrib.layers.xavier_initializer(dtype=dtype)
def make_conv_layer_weights(weights, scope, k, filters_in, filters_out, bias=True):
weights['{}/conv'.format(scope)] = tf.get_variable('{}/conv'.format(scope), [k, k, filters_in, filters_out], initializer=conv_initializer, dtype=dtype)
if bias:
weights['{}/bias'.format(scope)] = tf.get_variable('{}/bias'.format(scope), [filters_out], initializer=bias_initializer, dtype=dtype)
def make_fc_layer_weights(weights, scope, dims_in, dims_out):
weights['{}/fc'.format(scope)] = tf.get_variable('{}/fc'.format(scope), [dims_in, dims_out], initializer=fc_initializer, dtype=dtype)
weights['{}/bias'.format(scope)] = tf.get_variable('{}/bias'.format(scope), [dims_out], initializer=bias_initializer, dtype=dtype)
for block_name in self.blocks:
if block_name == 'input':
make_conv_layer_weights(weights, block_name, k=3, filters_in=self.channels, filters_out=64)
elif 'res' in block_name:
j = int(block_name[-1])
last_block_filter = 64 if j == 0 else 64 * 2 ** (j - 1)
this_block_filter = 64 if j == 0 else last_block_filter * 2
print(block_name, last_block_filter, this_block_filter)
make_conv_layer_weights(weights, '{}/shortcut'.format(block_name), k=1, filters_in=last_block_filter,
filters_out=this_block_filter, bias=False)
for i in range(self.num_parts_per_res_block):
make_conv_layer_weights(weights, '{}/part{}'.format(block_name, i), k=3,
filters_in=last_block_filter if i == 0 else this_block_filter,
filters_out=this_block_filter)
elif block_name == 'output':
make_fc_layer_weights(weights, block_name, dims_in=512, dims_out=self.dim_output_train)
return weights
def forward_resnet224(self, inp, weights, prefix, reuse=False):
inp = tf.reshape(inp, [-1, self.img_size, self.img_size, self.channels])
for block_name in self.blocks:
if block_name == 'input':
conv = weights['{}/conv'.format(block_name)]
bias = weights['{}/bias'.format(block_name)]
inp = tf.nn.conv2d(inp, filter=conv, strides=[1, 2, 2, 1], padding="SAME") + bias
elif 'res' in block_name:
shortcut = inp
conv = weights['{}/shortcut/conv'.format(block_name)]
shortcut = tf.nn.conv2d(input=shortcut, filter=conv, strides=[1, 1, 1, 1], padding="SAME")
for part in range(self.num_parts_per_res_block):
part_name = 'part{}'.format(part)
scope = '{}/{}'.format(block_name, part_name)
conv = weights['{}/{}/conv'.format(block_name, part_name)]
bias = weights['{}/{}/bias'.format(block_name, part_name)]
inp = bn_relu_conv_block(inp=inp, conv=conv, bias=bias, reuse=reuse, scope=scope)
inp = shortcut + inp
elif 'maxpool' in block_name:
inp = tf.nn.max_pool(inp, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
elif 'output' in block_name:
inp = tf.reduce_mean(inp, [1, 2])
fc = weights['{}/fc'.format(block_name)]
bias = weights['{}/bias'.format(block_name)]
inp = tf.matmul(inp, fc) + bias
if 'val' in prefix:
inp = tf.gather(inp, tf.range(self.dim_output_val), axis=1)
return inp
def forward_resnet84(self, inp, weights, prefix, reuse=False):
inp = tf.reshape(inp, [-1, self.img_size, self.img_size, self.channels])
for block_name in self.blocks:
if block_name == 'input':
conv = weights['{}/conv'.format(block_name)]
bias = weights['{}/bias'.format(block_name)]
inp = tf.nn.conv2d(inp, filter=conv, strides=[1, 1, 1, 1], padding="SAME") + bias
elif 'res' in block_name:
shortcut = inp
conv = weights['{}/shortcut/conv'.format(block_name)]
shortcut = tf.nn.conv2d(input=shortcut, filter=conv, strides=[1, 1, 1, 1], padding="SAME")
for part in range(self.num_parts_per_res_block):
part_name = 'part{}'.format(part)
scope = '{}/{}'.format(block_name, part_name)
conv = weights['{}/{}/conv'.format(block_name, part_name)]
bias = weights['{}/{}/bias'.format(block_name, part_name)]
inp = bn_relu_conv_block(inp=inp, conv=conv, bias=bias, reuse=reuse, scope=scope)
inp = shortcut + inp
elif 'maxpool' in block_name:
inp = tf.nn.max_pool(inp, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
elif 'output' in block_name:
inp = tf.reduce_mean(inp, [1, 2])
fc = weights['{}/fc'.format(block_name)]
bias = weights['{}/bias'.format(block_name)]
inp = tf.matmul(inp, fc) + bias
if 'val' in prefix:
inp = tf.gather(inp, tf.range(self.dim_output_val), axis=1)
return inp
def wrap(self, inp, weights, prefix, reuse=False, scope=''):
unused = self.forward_resnet(inp, weights, prefix, reuse=False)
return self.forward_resnet(inp, weights, prefix, reuse=True)
if __name__ == '__main__':
import ipdb
FLAGS = flags.FLAGS
## Dataset/method options
flags.DEFINE_string('dataset', 'omniglot', 'omniglot or mnist or miniimagenet or celeba')
flags.DEFINE_integer('num_encoding_dims', -1, 'of unsupervised representation learning method')
flags.DEFINE_string('encoder', 'acai', 'acai or bigan or deepcluster or infogan')
## Training options
flags.DEFINE_integer('metatrain_iterations', 30000, 'number of metatraining iterations.')
flags.DEFINE_integer('meta_batch_size', 8, 'number of tasks sampled per meta-update')
flags.DEFINE_float('meta_lr', 0.001, 'the base learning rate of the generator')
flags.DEFINE_float('update_lr', 0.05, 'step size alpha for inner gradient update.')
flags.DEFINE_integer('inner_update_batch_size_train', 1,
'number of examples used for inner gradient update (K for K-shot learning).')
flags.DEFINE_integer('inner_update_batch_size_val', 5, 'above but for meta-val')
flags.DEFINE_integer('outer_update_batch_size', 5, 'number of examples used for outer gradient update')
flags.DEFINE_integer('num_updates', 5, 'number of inner gradient updates during training.')
flags.DEFINE_string('mt_mode', 'gtgt', 'meta-training mode (for sampling, labeling): gtgt or encenc')
flags.DEFINE_string('mv_mode', 'gtgt', 'meta-validation mode (for sampling, labeling): gtgt or encenc')
flags.DEFINE_integer('num_classes_train', 5, 'number of classes used in classification for meta-training')
flags.DEFINE_integer('num_classes_val', 5, 'number of classes used in classification for meta-validation.')
flags.DEFINE_float('margin', 0.0, 'margin for generating partitions using random hyperplanes')
flags.DEFINE_integer('num_partitions', 1, 'number of partitions, -1 for same as number of meta-training tasks')
flags.DEFINE_string('partition_algorithm', 'kmeans', 'hyperplanes or kmeans')
flags.DEFINE_integer('num_clusters', -1, 'number of clusters for kmeans')
flags.DEFINE_boolean('scaled_encodings', True, 'if True, use randomly scaled encodings for kmeans')
flags.DEFINE_boolean('on_encodings', False, 'if True, train MAML on top of encodings')
flags.DEFINE_integer('num_hidden_layers', 2, 'number of mlp hidden layers')
flags.DEFINE_integer('num_parallel_calls', 8, 'for loading data')
flags.DEFINE_integer('gpu', 7, 'CUDA_VISIBLE_DEVICES=')
## Model options
flags.DEFINE_string('norm', 'batch_norm', 'batch_norm, layer_norm, or None')
flags.DEFINE_integer('num_filters', 32, 'number of filters for each conv layer')
flags.DEFINE_bool('conv', True, 'whether or not to use a convolutional network')
flags.DEFINE_bool('max_pool', False, 'Whether or not to use max pooling rather than strided convolutions')
flags.DEFINE_bool('stop_grad', False, 'if True, do not use second derivatives in meta-optimization (for speed)')
## Logging, saving, and testing options
flags.DEFINE_bool('log', True, 'if false, do not log summaries, for debugging code.')
flags.DEFINE_string('logdir', './log', 'directory for summaries and checkpoints.')
flags.DEFINE_bool('resume', True, 'resume training if there is a model available')
flags.DEFINE_bool('train', True, 'True to train, False to test.')
flags.DEFINE_integer('test_iter', -1, 'iteration to load model (-1 for latest model)')
flags.DEFINE_bool('test_set', False, 'Set to true to test on the the test set, False for the validation set.')
flags.DEFINE_integer('log_inner_update_batch_size_val', -1,
'specify log directory iubsv. (use to test with different iubsv)')
flags.DEFINE_float('train_update_lr', -1,
'value of inner gradient step step during training. (use if you want to test with a different value)')
flags.DEFINE_bool('save_checkpoints', False, 'if True, save model weights as checkpoints')
flags.DEFINE_bool('debug', False, 'if True, use tf debugger')
flags.DEFINE_string('suffix', '', 'suffix for an exp_string')
flags.DEFINE_bool('from_scratch', False, 'fast-adapt from scratch')
flags.DEFINE_integer('num_eval_tasks', 1000, 'number of tasks to meta-test on')
# Imagenet
flags.DEFINE_string('input_type', 'images_84x84',
'features or features_processed or images_fullsize or images_84x84')
flags.DEFINE_string('data_dir', '/data3/kylehsu/data', 'location of data')
flags.DEFINE_bool('resnet', False, 'use resnet architecture')
flags.DEFINE_integer('num_res_blocks', 5, 'number of resnet blocks')
flags.DEFINE_integer('num_parts_per_res_block', 2, 'number of bn-relu-conv parts in a res block')
FLAGS.resnet = True
maml = MAML(dim_input=3*84*84, dim_output_train=10, dim_output_val=5, test_num_updates=5)
maml.channels = 3
maml.img_size = 84
weights = maml.construct_resnet_weights()
input_ph = tf.placeholder(tf.float32)
unused = maml.forward_resnet(input_ph, weights, 'hi', reuse=False)
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
input = np.ones((1, 84 * 84 * 3), dtype=np.float32)
y = sess.run(maml.forward_resnet(input_ph, weights, 'val', reuse=True), {input_ph: input})
ipdb.set_trace()
x=1
