def __init__(self, log_dir, log_file='log.txt', exp_name=None):
"""
Initialize a Logger.
Args:
log_dir (string): A directory for saving results to. If
``None``, defaults to a temp directory of the form
``/tmp/experiments/somerandomnumber``.
log_file (string): Name for the tab-separated-value file
containing metrics logged throughout a training run.
Defaults to ``progress.txt``.
exp_name (string): Experiment name. If you run multiple training
runs and give them all the same ``exp_name``, the plotter
will know to group them. (Use case: if you run the same
hyperparameter configuration with multiple random seeds, you
should give them all the same ``exp_name``.)
"""
log_file = log_file if log_file.endswith('.txt') else log_file + '.txt'
self.log_dir = log_dir or f"/tmp/experiments/{time.time()}"
if osp.exists(self.log_dir):
print(
f"Warning: Log dir {self.log_dir} already exists! Storing info there anyway."
)
else:
os.makedirs(self.log_dir)
self.output_file = open(osp.join(self.log_dir, log_file), 'w')
atexit.register(self.output_file.close)
pwc("Logging data to %s" % self.output_file.name, 'green', bold=True)
self.first_row = True
self.log_headers = []
self.log_current_row = {}
self.exp_name = exp_name
def save(self, print_terminal_info=True):
if hasattr(self, 'saver'):
if print_terminal_info:
pwc('Model saved', 'magenta')
return self.saver.save(self.sess, self.model_file)
else:
# no intention to treat no saver as an error, just print a warning message
pwc('No saver is available', 'magenta')
def _conv_block(x, kernel, bias, norm, max_pool, block_idx):
if norm == 'batch_norm':
norm = lambda x: tf.layers.batch_normalization(x, training=True)
elif norm == 'layer_norm':
norm = tc.layers.layer_norm
elif norm is None:
pwc('No normalization is used')
else:
raise NotImplementedError
strides = [1, 1, 1, 1] if max_pool else [1, 2, 2, 1]
with tf.variable_scope(f'conv_block_{block_idx}'):
x = tf.nn.conv2d(x, kernel, strides, 'SAME') + bias
x = norm(x)
x = tf.nn.relu(x)
x = tf.nn.max_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], 'VALID')
return x
def restore(self, model_file=None):
"""
To restore a specific version of model, set filename to the model stored in saved_models
"""
self.model_file = model_file
if not hasattr(self, 'saver'):
self.saver = self._setup_saver()
try:
if os.path.isdir(self.model_file):
ckpt = tf.train.get_checkpoint_state(self.model_file)
self.saver.restore(self.sess, ckpt.model_checkpoint_path)
else:
self.saver.restore(self.sess, self.model_file)
except:
pwc(
f'Model {self.model_name}: No saved model for "{self.name}" is found. \nStart Training from Scratch!',
'magenta')
else:
pwc(
f"Model {self.model_name}: Params for {self.name} are restored.",
'magenta')
def evaluate(self, eval_image=False, train_time=None):
if self.log_tensorboard:
if train_time is None:
raise ValueError
summary = self.sess.run(
self.graph_summary,
feed_dict={self.image: self.data_generator.sample()})
self.writer.add_summary(summary, train_time)
if eval_image:
t, st_image = timeit(lambda: self.sess.run(
self.st_image, feed_dict={self.image: self.eval_image}))
pwc(f'Time taking to transfer an image: {t:.2f}s', color='green')
st_image = np.squeeze(st_image)
# get image path
image_filename, _ = osp.splitext(self.eval_image_path)
_, image_filename = osp.split(image_filename)
_, style_filename = osp.split(self.args['style_image_path'])
results_dir = self.args['results_dir']
if not osp.exists(results_dir):
os.mkdir(results_dir)
imsave(f'{results_dir}/{image_filename}-{style_filename}',
st_image,
format='jpg')
def display_var_info(vars, name='trainable'):
pwc(f'Print {name} variables', 'yellow')
count_params = 0
for v in vars:
name = v.name
if '/Adam' in name or 'beta1_power' in name or 'beta2_power' in name:
continue
v_params = np.prod(v.shape.as_list())
count_params += v_params
if '/b:' in name or '/biases' in name:
continue # Wx+b, bias is not interesting to look at => count params, but not print
pwc(f' {name}{" "*(100-len(name))} {v_params:d} params {v.shape}',
'yellow')
pwc(f'Total model parameters: {count_params*1e-6:0.2f} million', 'yellow')
def print_construction_complete(self):
pwc('Learner has been constructed.', color='cyan')
def make_data_tensor(self, training=True):
if training:
folders = self.metatrain_folders
num_total_batches = self.total_batch_num
else:
folders = self.metaval_folders
num_total_batches = 600
if training and os.path.exists('filelist.pkl'):
labels = np.arange(self.num_classes_per_task).repeat(self.num_samples_per_class).tolist()
with open('filelist.pkl', 'rb') as f:
all_filenames = pickle.load(f)
pwc('Load file names from filelist.pkl', 'magenta')
else:
all_filenames = []
for _ in tqdm.tqdm(range(num_total_batches), 'generating episodes'):
# sample data for each task
sampled_folders = random.sample(folders, self.num_classes_per_task)
random.shuffle(sampled_folders)
# for each task we have num_classes_per_task categories
# the same image may have different labels when it appears in different tasks
images_labels = get_images_labels(sampled_folders,
range(self.num_classes_per_task),
nb_samples=self.num_samples_per_class,
shuffle=False)
filenames, labels = list(zip(*images_labels))
all_filenames += filenames
if training:
with open('filelist.pkl', 'wb') as f:
pickle.dump(all_filenames, f)
pwc('Save all file names to filelist.pkl', 'magenta')
_, images = image_data`set(all_filenames, self.batch_size, self.image_size, norm_range=[0, 1], shuffle=False)
num_samples_per_task = self.num_classes_per_task * self.num_samples_per_class
all_image_batches, all_label_batches = [], []
# shuffle image&labels for each task
for i in range(self.num_tasks_per_batch):
# images for the current task
task_images = images[i * num_samples_per_task: (i + 1) * num_samples_per_task]
# as all labels of all task are the same, which is 0, 0, ..., 1, 1, ..., 2, 2, ..., 3, 3, ..., 4, 4, ...
label_batch = tf.convert_to_tensor(labels)
task_image_list, task_label_list = [], []
# shuffle [k, k+S, k+2*S, k+3*S, k+4*S]
for k in range(self.num_samples_per_class):
class_idxs = tf.range(0, self.num_classes_per_task)
class_idxs = tf.random_shuffle(class_idxs)
class_idxs = class_idxs * self.num_samples_per_class + k
task_image_list.append(tf.gather(task_images, class_idxs))
task_label_list.append(tf.gather(label_batch, class_idxs))
task_image_list = tf.concat(task_image_list, 0) # [C*S]
task_label_list = tf.concat(task_label_list, 0) # [C*S, 84, 84, 3]
all_image_batches.append(task_image_list)
all_label_batches.append(task_label_list)
# [4, C*S, 84*84*3]
all_image_batches = tf.stack(all_image_batches)
# [4, C*S]
all_label_batches = tf.stack(all_label_batches)
# [4, C*S, C]
all_label_batches = tf.one_hot(all_label_batches, self.num_classes_per_task)
pwc('Data pipeline has been constructed!', 'magenta')
return all_image_batches, all_label_batches
task_image_list = tf.concat(task_image_list, 0) # [C*S]
task_label_list = tf.concat(task_label_list, 0) # [C*S, 84, 84, 3]
all_image_batches.append(task_image_list)
all_label_batches.append(task_label_list)
# [4, C*S, 84*84*3]
all_image_batches = tf.stack(all_image_batches)
# [4, C*S]
all_label_batches = tf.stack(all_label_batches)
# [4, C*S, C]
all_label_batches = tf.one_hot(all_label_batches, self.num_classes_per_task)
pwc('Data pipeline has been constructed!', 'magenta')
return all_image_batches, all_label_batches
if __name__ == '__main__':
dg = DataGenerator(5, 2, 4, 500, 'dataset/miniimagenet/train', 'dataset/miniimagenet/test')
image, label = dg.make_data_tensor()
with tf.Session() as sess:
img, lb = sess.run([image, label])
pwc(f'{img.shape}')
lb = np.argmax(lb, 2)
print(lb)
pwc(f'{lb.shape}')
# comment out tf.one_hot above to enable the following test
for i in range(5):
save_image(img[i][lb[i]==1], f'./data_generator_test/img{i}.png')
def log(self, msg, color='green'):
"""Print a colorized message to stdout."""
pwc(msg, color, bold=True)
def __init__(self,
name,
args,
sess_config=None,
save=False,
log_tensorboard=False,
log_params=False,
log_stats=False,
device=None,
reuse=None,
graph=None):
""" Model, inherited from Module, further defines some boookkeeping functions,
such as session management, save & restore operations, tensorboard loggings, and etc.
Arguments:
name {str} -- Name of the model
args {dict} -- A dictionary which specifies necessary arguments for building graph
Keyword Arguments:
sess_config {tf.ConfigProto} -- session configuration (default: {None})
save {bool} -- Option for saving model (default: {True})
log_tensorboard {bool} -- Option for logging information to tensorboard (default: {False})
log_params {bool} -- Option for logging parameters to tensorboard (default: {False})
log_stats {bool} -- Option for logging score to tensorboard (default: {False})
device {[str or None]} -- Device where graph build upon {default: {None}}
"""
self.graph = graph if graph else tf.Graph()
# initialize session and global variables
if sess_config is None:
sess_config = tf.ConfigProto( #intra_op_parallelism_threads=2,
#inter_op_parallelism_threads=2,
allow_soft_placement=True)
# sess_config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
sess_config.gpu_options.allow_growth = True
self.sess = tf.Session(graph=self.graph, config=sess_config)
atexit.register(self.sess.close)
super().__init__(name,
args,
self.graph,
log_tensorboard=log_tensorboard,
log_params=log_params,
device=device,
reuse=reuse)
display_var_info(self.trainable_variables)
self.model_name = args['model_name']
if self.log_tensorboard:
self.graph_summary = self._setup_tensorboard_summary()
# rl-specific log configuration, not in self._build_graph to avoid being included in self.graph_summary
if log_stats:
self.logger = self._setup_logger(args['log_root_dir'],
self.model_name)
self.stats = self._setup_stats_logs(args['env_stats'])
if log_tensorboard or log_stats:
self.writer = self._setup_writer(args['log_root_dir'],
self.model_name)
self.sess.run(tf.variables_initializer(self.global_variables))
if save:
self.saver = self._setup_saver()
self.model_file = self._setup_model_path(args['model_root_dir'],
self.model_name)
pwc(f'{self.name} has been constructed', 'magenta')
def evaluate(self, n_iterations=1):
for i in range(n_iterations):
t, image = timeit(lambda: self.sess.run(self.gen_image))
pwc(f'\rEvaluation Time: {t}', 'magenta')
save_image(image, f'{self.results_dir}/eval_{i}.png')
def print_construction_complete(self):
pwc(f'{self.name} has been constructed', 'cyan')
