def main():
assert tf.executing_eagerly(), 'Only eager mode is supported.'
assert args.config is not None, 'You need to pass in model config file path'
assert args.data is not None, 'You need to pass in episode config file path'
assert args.env is not None, 'You need to pass in environ config file path'
assert args.tag is not None, 'You need to specify a tag'
log.info('Command line args {}'.format(args))
config = get_config(args.config, ExperimentConfig)
data_config = get_config(args.data, EpisodeConfig)
env_config = get_config(args.env, EnvironmentConfig)
log.info('Model: \n{}'.format(config))
log.info('Data episode: \n{}'.format(data_config))
log.info('Environment: \n{}'.format(env_config))
config.num_classes = data_config.nway # Assign num classes.
config.num_steps = data_config.maxlen
config.memory_net_config.max_classes = data_config.nway
config.memory_net_config.max_stages = data_config.nstage
config.memory_net_config.max_items = data_config.maxlen
config.oml_config.num_classes = data_config.nway
config.fix_unknown = data_config.fix_unknown # Assign fix unknown ID.
log.info('Number of classes {}'.format(data_config.nway))
if 'SLURM_JOB_ID' in os.environ:
log.info('SLURM job ID: {}'.format(os.environ['SLURM_JOB_ID']))
# Create save folder.
save_folder = os.path.join(env_config.results, env_config.dataset,
args.tag)
results_file = os.path.join(save_folder, 'results.pkl')
logfile = os.path.join(save_folder, 'results.csv')
# To get CNN features from.
model = build_pretrain_net(config)
if args.rnn:
rnn_memory = get_module("lstm",
"lstm",
model.backbone.get_output_dimension()[0],
config.lstm_config.hidden_dim,
layernorm=config.lstm_config.layernorm,
dtype=tf.float32)
memory = RNNEncoder(
"proto_plus_rnn_ssl_v4",
rnn_memory,
readout_type=config.mann_config.readout_type,
use_pred_beta_gamma=config.hybrid_config.use_pred_beta_gamma,
use_feature_fuse=config.hybrid_config.use_feature_fuse,
use_feature_fuse_gate=config.hybrid_config.use_feature_fuse_gate,
use_feature_scaling=config.hybrid_config.use_feature_scaling,
use_feature_memory_only=config.hybrid_config.
use_feature_memory_only,
skip_unk_memory_update=config.hybrid_config.skip_unk_memory_update,
use_ssl=config.hybrid_config.use_ssl,
use_ssl_beta_gamma_write=config.hybrid_config.
use_ssl_beta_gamma_write,
use_ssl_temp=config.hybrid_config.use_ssl_temp,
dtype=tf.float32)
rnn_model = RNNEncoderNet(config, model.backbone, memory)
f = lambda x: rnn_model.forward(x, is_training=False) # NOQA
else:
f = lambda x: model.backbone(x[0], is_training=False)[None, :, :
] # NOQA
K = config.num_classes
# Get dataset.
dataset = get_data_fs(env_config, load_train=True)
# Get data iterators.
if env_config.dataset in ["matterport"]:
data = get_dataiter_sim(
dataset,
data_config,
batch_size=config.optimizer_config.batch_size,
nchw=model.backbone.config.data_format == 'NCHW',
seed=args.seed)
else:
data = get_dataiter_continual(
dataset,
data_config,
batch_size=config.optimizer_config.batch_size,
nchw=model.backbone.config.data_format == 'NCHW',
save_additional_info=True,
random_box=data_config.random_box,
seed=args.seed)
if os.path.exists(results_file) and args.reeval:
# Re-display results.
results_all = pkl.load(open(results_file, 'rb'))
for split, name in zip(['trainval_fs', 'val_fs', 'test_fs'],
['Train', 'Val', 'Test']):
stats = get_stats(results_all[split], tmax=data_config.maxlen)
log_results(stats, prefix=name, filename=logfile)
else:
latest = latest_file(save_folder, 'weights-')
if args.rnn:
rnn_model.load(latest)
else:
model.load(latest)
data['trainval_fs'].reset()
data['val_fs'].reset()
data['test_fs'].reset()
results_all = {}
if args.testonly:
split_list = ['test_fs']
name_list = ['Test']
nepisode_list = [config.num_episodes]
else:
split_list = ['trainval_fs', 'val_fs', 'test_fs']
name_list = ['Train', 'Val', 'Test']
nepisode_list = [600, config.num_episodes, config.num_episodes]
for split, name, N in zip(split_list, name_list, nepisode_list):
data[split].reset()
r1 = evaluate(f, K, data[split], N)
stats = get_stats(r1, tmax=data_config.maxlen)
log_results(stats, prefix=name, filename=logfile)
results_all[split] = r1
pkl.dump(results_all, open(results_file, 'wb'))
def train(model,
dataiter,
dataiter_traintest,
dataiter_test,
ckpt_folder,
final_save_folder=None,
nshot_max=5,
maxlen=40,
logger=None,
writer=None,
in_stage=False,
is_chief=True,
reload_flag=None):
"""Trains the online few-shot model.
Args:
model: Model instance.
dataiter: Dataset iterator.
dataiter_test: Dataset iterator for validation.
save_folder: Path to save the checkpoints.
"""
N = model.max_train_steps
config = model.config.train_config
def try_log(*args, **kwargs):
if logger is not None:
logger.log(*args, **kwargs)
def try_flush(*args, **kwargs):
if logger is not None:
logger.flush()
r = None
keep = True
restart = False
best_val = 0.0
while keep:
keep = False
start = model.step.numpy()
if start > 0:
log.info('Restore from step {}'.format(start))
it = six.moves.xrange(start, N)
if is_chief:
it = tqdm(it, ncols=0)
for i, batch in zip(it, dataiter):
tf.summary.experimental.set_step(i + 1)
x = batch['x_s']
y = batch['y_s']
y = batch['y_s']
y_gt = batch['y_gt']
kwargs = {'y_gt': y_gt, 'flag': batch['flag_s']}
kwargs['writer'] = writer
loss = model.train_step(x, y, **kwargs)
if i == start and reload_flag is not None and not restart:
model.load(reload_flag, load_optimizer=True)
# Synchronize distributed weights.
if i == start and model._distributed and not restart:
import horovod.tensorflow as hvd
hvd.broadcast_variables(model.var_to_optimize(), root_rank=0)
hvd.broadcast_variables(model.optimizer.variables(),
root_rank=0)
if model.config.set_backbone_lr:
hvd.broadcast_variables(model._bb_optimizer.variables(),
root_rank=0)
if loss > 10.0 and i > start + config.steps_per_save:
# Something wrong happened.
log.error('Something wrong happened. loss = {}'.format(loss))
import pickle as pkl
pkl.dump(
batch,
open(os.path.join(final_save_folder, 'debug.pkl'), 'wb'))
log.error('debug file dumped')
restart = True
keep = True
latest = latest_file(ckpt_folder, 'weights-')
model.load(latest_file)
log.error(
'Reloaded latest checkpoint from {}'.format(latest_file))
break
# Evaluate.
if is_chief and ((i + 1) % config.steps_per_val == 0 or i == 0):
for key, data_it_ in zip(['train', 'val'],
[dataiter_traintest, dataiter_test]):
data_it_.reset()
r1 = evaluate(model, data_it_, 120)
r = get_stats(r1, nshot_max=nshot_max, tmax=maxlen)
for s in range(nshot_max):
try_log('online fs acc {}/s{}'.format(key, s), i + 1,
r['acc_nshot'][s] * 100.0)
try_log('online fs ap {}'.format(key), i + 1,
r['ap'] * 100.0)
try_log('lr', i + 1, model.learn_rate())
print()
# Save.
if is_chief and ((i + 1) % config.steps_per_save == 0 or i == 0):
model.save(
os.path.join(ckpt_folder, 'weights-{}'.format(i + 1)))
# Save the best checkpoint.
if r is not None:
if r['ap'] > best_val:
model.save(
os.path.join(final_save_folder,
'best-{}'.format(i + 1)))
best_val = r['ap']
# Write logs.
if is_chief and ((i + 1) % config.steps_per_log == 0 or i == 0):
try_log('loss', i + 1, loss)
try_flush()
# Update progress bar.
post_fix_dict = {}
post_fix_dict['lr'] = '{:.3e}'.format(model.learn_rate())
post_fix_dict['loss'] = '{:.3e}'.format(loss)
if r is not None:
post_fix_dict['ap_val'] = '{:.3f}'.format(r['ap'] * 100.0)
it.set_postfix(**post_fix_dict)
# Save.
if is_chief and final_save_folder is not None:
model.save(os.path.join(final_save_folder, 'weights-{}'.format(N)))
def main():
assert tf.executing_eagerly(), 'Only eager mode is supported.'
assert args.config is not None, 'You need to pass in model config file path'
assert args.data is not None, 'You need to pass in episode config file path'
assert args.env is not None, 'You need to pass in environ config file path'
assert args.tag is not None, 'You need to specify a tag'
log.info('Command line args {}'.format(args))
config = get_config(args.config, ExperimentConfig)
data_config = get_config(args.data, EpisodeConfig)
env_config = get_config(args.env, EnvironmentConfig)
log.info('Model: \n{}'.format(config))
log.info('Data episode: \n{}'.format(data_config))
log.info('Environment: \n{}'.format(env_config))
config.num_classes = data_config.nway # Assign num classes.
config.num_steps = data_config.maxlen
config.memory_net_config.max_classes = data_config.nway
config.memory_net_config.max_stages = data_config.nstage
config.memory_net_config.max_items = data_config.maxlen
config.oml_config.num_classes = data_config.nway
config.fix_unknown = data_config.fix_unknown # Assign fix unknown ID.
log.info('Number of classes {}'.format(data_config.nway))
if 'SLURM_JOB_ID' in os.environ:
log.info('SLURM job ID: {}'.format(os.environ['SLURM_JOB_ID']))
# Create save folder.
save_folder = os.path.join(env_config.results, env_config.dataset,
args.tag)
if not args.reeval:
model = build_pretrain_net(config)
mem_model = build_net(config, backbone=model.backbone)
reload_flag = None
restore_steps = 0
# Checkpoint folder.
ckpt_path = env_config.checkpoint
if len(ckpt_path) > 0 and os.path.exists(ckpt_path):
ckpt_folder = os.path.join(ckpt_path, os.environ['SLURM_JOB_ID'])
else:
ckpt_folder = save_folder
# Reload previous checkpoint.
if os.path.exists(ckpt_folder) and not args.eval:
latest = latest_file(ckpt_folder, 'weights-')
if latest is not None:
log.info('Checkpoint already exists. Loading from {}'.format(
latest))
mem_model.load(latest) # Not loading optimizer weights here.
reload_flag = latest
restore_steps = int(reload_flag.split('-')[-1])
if not args.eval:
save_config(config, save_folder)
# Create TB logger.
if not args.eval:
writer = tf.summary.create_file_writer(save_folder)
logger = ExperimentLogger(writer)
# Get dataset.
dataset = get_data_fs(env_config, load_train=True)
# Get data iterators.
if env_config.dataset in ["matterport"]:
data = get_dataiter_sim(
dataset,
data_config,
batch_size=config.optimizer_config.batch_size,
nchw=mem_model.backbone.config.data_format == 'NCHW',
seed=args.seed + restore_steps)
else:
data = get_dataiter_continual(
dataset,
data_config,
batch_size=config.optimizer_config.batch_size,
nchw=mem_model.backbone.config.data_format == 'NCHW',
save_additional_info=True,
random_box=data_config.random_box,
seed=args.seed + restore_steps)
# Load model, training loop.
if not args.eval:
if args.pretrain is not None and reload_flag is None:
model.load(latest_file(args.pretrain, 'weights-'))
if config.freeze_backbone:
model.set_trainable(False) # Freeze the network.
log.info('Backbone network is now frozen')
with writer.as_default() if writer is not None else dummy_context_mgr(
) as gs:
train(mem_model,
data['train_fs'],
data['trainval_fs'],
data['val_fs'],
ckpt_folder,
final_save_folder=save_folder,
maxlen=data_config.maxlen,
logger=logger,
writer=writer,
in_stage=config.in_stage,
reload_flag=reload_flag)
else:
results_file = os.path.join(save_folder, 'results.pkl')
logfile = os.path.join(save_folder, 'results.csv')
if os.path.exists(results_file) and args.reeval:
# Re-display results.
results_all = pkl.load(open(results_file, 'rb'))
for split, name in zip(['trainval_fs', 'val_fs', 'test_fs'],
['Train', 'Val', 'Test']):
stats = get_stats(results_all[split], tmax=data_config.maxlen)
log_results(stats, prefix=name, filename=logfile)
else:
# Load the most recent checkpoint.
if args.usebest:
latest = latest_file(save_folder, 'best-')
else:
latest = latest_file(save_folder, 'weights-')
if latest is not None:
mem_model.load(latest)
else:
latest = latest_file(args.pretrain, 'weights-')
if latest is not None:
mem_model.load(latest)
else:
raise ValueError('Checkpoint not found')
data['trainval_fs'].reset()
data['val_fs'].reset()
data['test_fs'].reset()
results_all = {}
if args.testonly:
split_list = ['test_fs']
name_list = ['Test']
nepisode_list = [config.num_episodes]
else:
split_list = ['trainval_fs', 'val_fs', 'test_fs']
name_list = ['Train', 'Val', 'Test']
nepisode_list = [600, config.num_episodes, config.num_episodes]
for split, name, N in zip(split_list, name_list, nepisode_list):
data[split].reset()
r1 = evaluate(mem_model, data[split], N, verbose=True)
stats = get_stats(r1, tmax=data_config.maxlen)
log_results(stats, prefix=name, filename=logfile)
results_all[split] = r1
pkl.dump(results_all, open(results_file, 'wb'))
def test_load_one(folder, seed=0):
# log.info('Command line args {}'.format(args))
config_file = os.path.join(folder, 'config.prototxt')
config = get_config(config_file, ExperimentConfig)
# config.c4_config.data_format = 'NHWC'
# config.resnet_config.data_format = 'NHWC'
if 'omniglot' in folder:
if 'ssl' in folder:
data_config_file = 'configs/episodes/roaming-omniglot/roaming-omniglot-150-ssl.prototxt' # NOQA
else:
data_config_file = 'configs/episodes/roaming-omniglot/roaming-omniglot-150.prototxt' # NOQA
env_config_file = 'configs/environ/roaming-omniglot-docker.prototxt'
elif 'rooms' in folder:
if 'ssl' in folder:
data_config_file = 'configs/episodes/roaming-rooms/roaming-rooms-100.prototxt' # NOQA
else:
data_config_file = 'configs/episodes/roaming-rooms/romaing-rooms-100.prototxt' # NOQA
env_config_file = 'configs/environ/roaming-rooms-docker.prototxt'
data_config = get_config(data_config_file, EpisodeConfig)
env_config = get_config(env_config_file, EnvironmentConfig)
log.info('Model: \n{}'.format(config))
log.info('Data episode: \n{}'.format(data_config))
log.info('Environment: \n{}'.format(env_config))
config.num_classes = data_config.nway # Assign num classes.
config.num_steps = data_config.maxlen
config.memory_net_config.max_classes = data_config.nway
config.memory_net_config.max_stages = data_config.nstage
config.memory_net_config.max_items = data_config.maxlen
config.oml_config.num_classes = data_config.nway
config.fix_unknown = data_config.fix_unknown # Assign fix unknown ID.
log.info('Number of classes {}'.format(data_config.nway))
model = build_pretrain_net(config)
mem_model = build_net(config, backbone=model.backbone)
reload_flag = None
restore_steps = 0
# Reload previous checkpoint.
latest = latest_file(folder, 'best-')
if latest is None:
latest = latest_file(folder, 'weights-')
assert latest is not None, "Checkpoint not found."
if latest is not None:
log.info('Checkpoint already exists. Loading from {}'.format(latest))
mem_model.load(latest) # Not loading optimizer weights here.
reload_flag = latest
restore_steps = int(reload_flag.split('-')[-1])
# Get dataset.
dataset = get_data_fs(env_config, load_train=True)
# Get data iterators.
if env_config.dataset in ["roaming-rooms", "matterport"]:
data = get_dataiter_sim(
dataset,
data_config,
batch_size=config.optimizer_config.batch_size,
nchw=mem_model.backbone.config.data_format == 'NCHW',
seed=seed + restore_steps)
else:
data = get_dataiter_continual(
dataset,
data_config,
batch_size=config.optimizer_config.batch_size,
nchw=mem_model.backbone.config.data_format == 'NCHW',
save_additional_info=True,
random_box=data_config.random_box,
seed=seed + restore_steps)
# Load the most recent checkpoint.
latest = latest_file(folder, 'best-')
if latest is None:
latest = latest_file(folder, 'weights-')
data['trainval_fs'].reset()
data['val_fs'].reset()
data['test_fs'].reset()
results_all = {}
split_list = ['trainval_fs', 'val_fs', 'test_fs']
name_list = ['Train', 'Val', 'Test']
nepisode_list = [5, 5, 5]
# nepisode_list = [600, config.num_episodes, config.num_episodes]
for split, name, N in zip(split_list, name_list, nepisode_list):
# print(name)
data[split].reset()
r1 = evaluate(mem_model, data[split], N, verbose=False)
stats = get_stats(r1, tmax=data_config.maxlen)
print(split, stats['ap'])