def test_fetch_dataloaders(self):
dataloaders = fetch_dataloaders(['train', 'test'], self.task)
dl_train = dataloaders['train']
dl_test = dataloaders['test']
N = self.num_classes * self.num_samples
T = self.num_classes * self.num_query
for i, (sup_samples, sup_labels) in enumerate(dl_train):
self.assertTrue(sup_samples.size() == (N, 1, 28, 28))
self.assertTrue(i == 0)
for i, (que_samples, que_labels) in enumerate(dl_test):
self.assertTrue(que_samples.size() == (T, 1, 28, 28))
self.assertTrue(i == 0)
def evaluate(model, loss_fn, meta_classes, task_type, metrics, params, split):
"""
Evaluate the model on `num_steps` batches.
Args:
model: TPN model
loss_fn: a loss function
meta_classes: (list) a list of classes to be evaluated in meta-training or meta-testing
task_type: (subclass of FewShotTask) a type for generating tasks
metrics: (dict) a dictionary of functions that compute a metric using
the output and labels of each batch
params: (Params) hyperparameters
split: (string) 'train' if evaluate on 'meta-training' and
'val' if evaluate on 'meta-validating' and
'test' if evaluate on 'meta-testing'
"""
# params information
SEED = params.SEED
num_classes = params.num_classes
num_samples = params.num_samples
num_query = params.num_query
num_steps = params.num_steps
# set model to evaluation mode
model.eval()
# summary for current eval loop
summ = []
# compute metrics over the dataset
for episode in range(num_steps):
# Make a single task
# Make dataloaders to load support set and query set
task = task_type(meta_classes, num_classes, num_samples, num_query)
dataloaders = fetch_dataloaders(['train', 'test'], task)
dl_sup = dataloaders['train']
dl_que = dataloaders['test']
X_sup, Y_sup = dl_sup.__iter__().next()
X_que, Y_que = dl_que.__iter__().next()
# move to GPU if available
if params.cuda:
X_sup, Y_sup = X_sup.cuda(async=True), Y_sup.cuda(async=True)
X_que, Y_que = X_que.cuda(async=True), Y_que.cuda(async=True)
# Evaluate the model given a task
Y_que_hat = model(X_sup, Y_sup, X_que)
loss = loss_fn(Y_que_hat, Y_que)
# extract data from torch Variable, move to cpu, convert to numpy arrays
Y_que_hat = Y_que_hat.data.cpu().numpy()
Y_que = Y_que.data.cpu().numpy()
# compute all metrics on this batch
summary_batch = {
metric: metrics[metric](Y_que_hat, Y_que)
for metric in metrics
}
summary_batch['loss'] = loss.item()
summ.append(summary_batch)
# compute mean of all metrics in summary
metrics_mean = {
metric: np.mean([x[metric] for x in summ])
for metric in summ[0]
}
metrics_string = " ; ".join("{}: {:05.3f}".format(k, v)
for k, v in metrics_mean.items())
logging.info("- [" + split.upper() + "] Eval metrics : " + metrics_string)
return metrics_mean
def train_and_evaluate(model,
meta_train_classes,
meta_val_classes,
meta_test_classes,
task_type,
optimizer,
scheduler,
loss_fn,
metrics,
params,
model_dir,
restore_file=None):
"""
Train the model and evaluate every `save_summary_steps`.
Args:
model: TPN model
meta_train_classes: (list) the classes for meta-training
meta_val_classes: (list) the classes for meta-validating
meta_test_classes: (list) the classes for meta-testing
task_type: (subclass of FewShotTask) a type for generating tasks
optimizer: (torch.optim) optimizer for parameters of model
scheduler: (torch.optim.lr_scheduler) scheduler for decaying learning rate
loss_fn: a loss function
metrics: (dict) a dictionary of functions that compute a metric using
the output and labels of each batch
params: (Params) hyperparameters
model_dir: (string) directory containing config, weights and log
restore_file: (string) optional- name of file to restore from
(without its extension .pth.tar)
"""
# reload weights from restore_file if specified
if restore_file is not None:
restore_path = os.path.join(args.model_dir,
args.restore_file + '.pth.tar')
logging.info("Restoring parameters from {}".format(restore_path))
utils.load_checkpoint(restore_path, model, optimizer)
# params information
num_classes = params.num_classes
num_samples = params.num_samples
num_query = params.num_query
# validation accuracy
best_val_loss = float('inf')
# For plotting to see summerized training procedure
plot_history = {
'train_loss': [],
'train_acc': [],
'val_loss': [],
'val_acc': [],
'test_loss': [],
'test_acc': []
}
with tqdm(total=params.num_episodes) as t:
for episode in range(params.num_episodes):
# Run one episode
logging.info("Episode {}/{}".format(episode + 1,
params.num_episodes))
scheduler.step()
# Train a model on a single task (episode).
# TODO meta-batch of tasks
task = task_type(meta_train_classes, num_classes, num_samples,
num_query)
dataloaders = fetch_dataloaders(['train', 'test'], task)
_ = train_single_task(model, optimizer, loss_fn, dataloaders,
metrics, params)
# print(episode, _)
# Evaluate on train, val, test dataset given a number of tasks (params.num_steps)
if (episode + 1) % params.save_summary_steps == 0:
train_metrics = evaluate(model, loss_fn, meta_train_classes,
task_type, metrics, params, 'train')
val_metrics = evaluate(model, loss_fn, meta_val_classes,
task_type, metrics, params, 'val')
test_metrics = evaluate(model, loss_fn, meta_test_classes,
task_type, metrics, params, 'test')
train_loss = train_metrics['loss']
val_loss = val_metrics['loss']
test_loss = test_metrics['loss']
train_acc = train_metrics['accuracy']
val_acc = val_metrics['accuracy']
test_acc = test_metrics['accuracy']
is_best = val_loss <= best_val_loss
# Save weights
utils.save_checkpoint({
'episode': episode + 1,
'state_dict': model.state_dict(),
'optim_dict': optimizer.state_dict()
},
is_best=is_best,
checkpoint=model_dir)
# If best_test, best_save_path
if is_best:
logging.info("- Found new best accuracy")
best_val_loss = val_loss
# Save best test metrics in a json file in the model directory
best_train_json_path = os.path.join(
model_dir, "metrics_train_best_weights.json")
utils.save_dict_to_json(train_metrics,
best_train_json_path)
best_val_json_path = os.path.join(
model_dir, "metrics_val_best_weights.json")
utils.save_dict_to_json(val_metrics, best_val_json_path)
best_test_json_path = os.path.join(
model_dir, "metrics_test_best_weights.json")
utils.save_dict_to_json(test_metrics, best_test_json_path)
# Save latest test metrics in a json file in the model directory
last_train_json_path = os.path.join(
model_dir, "metrics_train_last_weights.json")
utils.save_dict_to_json(train_metrics, last_train_json_path)
last_val_json_path = os.path.join(
model_dir, "metrics_val_last_weights.json")
utils.save_dict_to_json(val_metrics, last_val_json_path)
last_test_json_path = os.path.join(
model_dir, "metrics_test_last_weights.json")
utils.save_dict_to_json(test_metrics, last_test_json_path)
plot_history['train_loss'].append(train_loss)
plot_history['train_acc'].append(train_acc)
plot_history['val_loss'].append(val_loss)
plot_history['val_acc'].append(val_acc)
plot_history['test_loss'].append(test_loss)
plot_history['test_acc'].append(test_acc)
utils.plot_training_results(args.model_dir, plot_history)
t.set_postfix(
tr_acc='{:05.3f}'.format(train_acc),
te_acc='{:05.3f}'.format(test_acc),
tr_loss='{:05.3f}'.format(train_loss),
te_loss='{:05.3f}'.format(test_loss))
print('\n')
t.update()
def evaluate(model, loss_fn, meta_classes, task_lr, task_type, metrics, params,
args, split):
"""
Evaluate the model on `num_steps` batches.
Args:
model: (MetaLearner) a meta-learner that is trained on MAML
loss_fn: a loss function
meta_classes: (list) a list of classes to be evaluated in meta-training or meta-testing
task_lr: (float) a task-specific learning rate
task_type: (subclass of FewShotTask) a type for generating tasks
metrics: (dict) a dictionary of functions that compute a metric using
the output and labels of each batch
params: (Params) hyperparameters
split: (string) 'train' if evaluate on 'meta-training' and
'test' if evaluate on 'meta-testing' TODO 'meta-validating'
"""
# params information
SEED = params.SEED
num_classes = params.num_classes
num_samples = params.num_samples
num_query = params.num_query
num_steps = params.num_steps
num_eval_updates = params.num_eval_updates
# set model to evaluation mode
# NOTE eval() is not needed since everytime task is varying and batchnorm
# should compute statistics within the task.
# model.eval()
# summary for current eval loop
summ = []
# compute metrics over the dataset
for episode in range(num_steps):
# Make a single task
# Make dataloaders to load support set and query set
task = task_type(args, meta_classes, num_classes, num_samples,
num_query)
dataloaders = fetch_dataloaders(['train', 'test'], task, args.pkl_path,
params)
dl_sup = dataloaders['train']
dl_que = dataloaders['test']
X_sup, Y_sup = dl_sup.__iter__().next()
X_que, Y_que = dl_que.__iter__().next()
#print ( Y_que.detach().numpy() )
# move to GPU if available
if params.cuda:
X_sup, Y_sup = X_sup.cuda(), Y_sup.cuda()
X_que, Y_que = X_que.cuda(), Y_que.cuda()
# Direct optimization
net_clone = copy.deepcopy(model)
optim = torch.optim.SGD(net_clone.parameters(), lr=task_lr)
for _ in range(num_eval_updates):
Y_sup_hat = net_clone(X_sup)
loss = loss_fn(Y_sup_hat, Y_sup)
optim.zero_grad()
loss.backward()
optim.step()
Y_que_hat = net_clone(X_que)
loss = loss_fn(Y_que_hat, Y_que)
# # clear previous gradients, compute gradients of all variables wrt loss
# def zero_grad(params):
# for p in params:
# if p.grad is not None:
# p.grad.zero_()
# # NOTE In Meta-SGD paper, num_eval_updates=1 is enough
# for _ in range(num_eval_updates):
# Y_sup_hat = model(X_sup)
# loss = loss_fn(Y_sup_hat, Y_sup)
# zero_grad(model.parameters())
# grads = torch.autograd.grad(loss, model.parameters())
# # step() manually
# adapted_state_dict = model.cloned_state_dict()
# adapted_params = OrderedDict()
# for (key, val), grad in zip(model.named_parameters(), grads):
# adapted_params[key] = val - task_lr * grad
# adapted_state_dict[key] = adapted_params[key]
# Y_que_hat = model(X_que, adapted_state_dict)
# loss = loss_fn(Y_que_hat, Y_que) # NOTE !!!!!!!!
# extract data from torch Variable, move to cpu, convert to numpy arrays
Y_que_hat = Y_que_hat.data.cpu().numpy()
Y_que = Y_que.data.cpu().numpy()
# compute all metrics on this batch
summary_batch = {
metric: metrics[metric](Y_que_hat, Y_que)
for metric in metrics
}
summary_batch['loss'] = loss.item()
summ.append(summary_batch)
# compute mean of all metrics in summary
metrics_mean = {
metric: np.mean([x[metric] for x in summ])
for metric in summ[0]
}
metrics_string = " ; ".join("{}: {:05.6f}".format(k, v)
for k, v in metrics_mean.items())
logging.info("- [" + split.upper() + "] Eval metrics : " + metrics_string)
print(metrics_string)
return metrics_mean
def train_and_evaluate(model,
meta_train_classes,
meta_test_classes,
task_type,
meta_optimizer,
loss_fn,
metrics,
params,
model_dir,
restore_file=None):
"""
Train the model and evaluate every `save_summary_steps`.
Args:
model: (MetaLearner) a meta-learner for MAML algorithm
meta_train_classes: (list) the classes for meta-training
meta_train_classes: (list) the classes for meta-testing
task_type: (subclass of FewShotTask) a type for generating tasks
meta_optimizer: (torch.optim) an meta-optimizer for MetaLearner
loss_fn: a loss function
metrics: (dict) a dictionary of functions that compute a metric using
the output and labels of each batch
params: (Params) hyperparameters
model_dir: (string) directory containing config, weights and log
restore_file: (string) optional- name of file to restore from
(without its extension .pth.tar)
TODO Validation classes
"""
# reload weights from restore_file if specified
if restore_file is not None:
restore_path = os.path.join(args.model_dir,
args.restore_file + '.pth.tar')
logging.info("Restoring parameters from {}".format(restore_path))
utils.load_checkpoint(restore_path, model, meta_optimizer)
# params information
num_classes = params.num_classes
num_samples = params.num_samples
num_query = params.num_query
num_inner_tasks = params.num_inner_tasks
meta_lr = params.meta_lr
# TODO validation accuracy
best_test_acc = 0.0
# For plotting to see summerized training procedure
plot_history = {
'train_loss': [],
'train_acc': [],
'test_loss': [],
'test_acc': []
}
with tqdm(total=params.num_episodes) as t:
for episode in range(params.num_episodes):
# Run one episode
logging.info("Episode {}/{}".format(episode + 1,
params.num_episodes))
# Run inner loops to get adapted parameters (theta_t`)
adapted_state_dicts = []
dataloaders_list = []
for n_task in range(num_inner_tasks):
task = task_type(meta_train_classes, num_classes, num_samples,
num_query)
dataloaders = fetch_dataloaders(['train', 'test', 'meta'],
task)
# Perform a gradient descent to meta-learner on the task
a_dict = train_single_task(model, loss_fn, dataloaders, params)
# Store adapted parameters
# Store dataloaders for meta-update and evaluation
adapted_state_dicts.append(a_dict)
dataloaders_list.append(dataloaders)
# Update the parameters of meta-learner
# Compute losses with adapted parameters along with corresponding tasks
# Updated the parameters of meta-learner using sum of the losses
meta_loss = 0
for n_task in range(num_inner_tasks):
dataloaders = dataloaders_list[n_task]
dl_meta = dataloaders['meta']
X_meta, Y_meta = dl_meta.__iter__().next()
if params.cuda:
X_meta, Y_meta = X_meta.cuda(async=True), Y_meta.cuda(
async=True)
a_dict = adapted_state_dicts[n_task]
Y_meta_hat = model(X_meta, a_dict)
loss_t = loss_fn(Y_meta_hat, Y_meta)
meta_loss += loss_t
meta_loss /= float(num_inner_tasks)
# print(meta_loss.item())
# Meta-update using meta_optimizer
meta_optimizer.zero_grad()
meta_loss.backward()
meta_optimizer.step()
# print(model.task_lr.values())
# Evaluate model on new task
# Evaluate on train and test dataset given a number of tasks (params.num_steps)
if (episode + 1) % params.save_summary_steps == 0:
train_metrics = evaluate(model, loss_fn, meta_train_classes,
task_type, metrics, params, 'train')
test_metrics = evaluate(model, loss_fn, meta_test_classes,
task_type, metrics, params, 'test')
train_loss = train_metrics['loss']
test_loss = test_metrics['loss']
train_acc = train_metrics['accuracy']
test_acc = test_metrics['accuracy']
is_best = test_acc >= best_test_acc
# Save weights
utils.save_checkpoint(
{
'episode': episode + 1,
'state_dict': model.state_dict(),
'optim_dict': meta_optimizer.state_dict()
},
is_best=is_best,
checkpoint=model_dir)
# If best_test, best_save_path
if is_best:
logging.info("- Found new best accuracy")
best_test_acc = test_acc
# Save best test metrics in a json file in the model directory
best_train_json_path = os.path.join(
model_dir, "metrics_train_best_weights.json")
utils.save_dict_to_json(train_metrics,
best_train_json_path)
best_test_json_path = os.path.join(
model_dir, "metrics_test_best_weights.json")
utils.save_dict_to_json(test_metrics, best_test_json_path)
# Save latest test metrics in a json file in the model directory
last_train_json_path = os.path.join(
model_dir, "metrics_train_last_weights.json")
utils.save_dict_to_json(train_metrics, last_train_json_path)
last_test_json_path = os.path.join(
model_dir, "metrics_test_last_weights.json")
utils.save_dict_to_json(test_metrics, last_test_json_path)
plot_history['train_loss'].append(train_loss)
plot_history['train_acc'].append(train_acc)
plot_history['test_loss'].append(test_loss)
plot_history['test_acc'].append(test_acc)
t.set_postfix(tr_acc='{:05.3f}'.format(train_acc),
te_acc='{:05.3f}'.format(test_acc),
tr_loss='{:05.3f}'.format(train_loss),
te_loss='{:05.3f}'.format(test_loss))
print('\n')
t.update()
utils.plot_training_results(args.model_dir, plot_history)
def evaluate(model, loss_fn, meta_classes, task_type, metrics, params, split):
"""
Evaluate the model on `num_steps` batches.
Args:
model: (MetaLearner) a meta-learner that is trained on MAML
loss_fn: a loss function
meta_classes: (list) a list of classes to be evaluated in meta-training or meta-testing
task_type: (subclass of FewShotTask) a type for generating tasks
metrics: (dict) a dictionary of functions that compute a metric using
the output and labels of each batch
params: (Params) hyperparameters
split: (string) 'train' if evaluate on 'meta-training' and
'val' if evaluate on 'meta-validating' and
'test' if evaluate on 'meta-testing'
"""
# set model to evaluation mode
# NOTE eval() is not needed since everytime task is varying and batchnorm
# should compute statistics within the task.
# model.eval()
# summary for current eval loop
summ = []
# compute metrics over the dataset
for episode in range(params.num_steps):
# Make a single task
task = task_type(meta_classes, params.num_classes, params.num_samples,
params.num_query)
dataloaders = fetch_dataloaders(['train', 'test'], task)
dl_sup = dataloaders['train']
dl_que = dataloaders['test']
X_sup, Y_sup = dl_sup.__iter__().next()
X_que, Y_que = dl_que.__iter__().next()
# move to GPU if available
if params.cuda:
X_sup, Y_sup = X_sup.cuda(async=True), Y_sup.cuda(async=True)
X_que, Y_que = X_que.cuda(async=True), Y_que.cuda(async=True)
# clear previous gradients, compute gradients of all variables wrt loss
def zero_grad(params):
for p in params:
if p.grad is not None:
p.grad.zero_()
# Single inner gradient update
Y_sup_hat = model(X_sup)
loss = loss_fn(Y_sup_hat, Y_sup)
zero_grad(model.parameters())
grads = torch.autograd.grad(loss, model.parameters())
# step() manually
adapted_state_dict = model.cloned_state_dict()
adapted_params = OrderedDict()
for (key, val), grad in zip(model.named_parameters(), grads):
# NOTE Here Meta-SGD is different from naive MAML
task_lr = model.task_lr[key]
adapted_params[key] = val - task_lr * grad
adapted_state_dict[key] = adapted_params[key]
Y_que_hat = model(X_que, adapted_state_dict)
loss = loss_fn(Y_que_hat, Y_que) # NOTE !!!!!!!!
# extract data from torch Variable, move to cpu, convert to numpy arrays
Y_que_hat = Y_que_hat.data.cpu().numpy()
Y_que = Y_que.data.cpu().numpy()
# compute all metrics on this batch
summary_batch = {
metric: metrics[metric](Y_que_hat, Y_que)
for metric in metrics
}
summary_batch['loss'] = loss.item()
summ.append(summary_batch)
# compute mean of all metrics in summary
metrics_mean = {
metric: np.mean([x[metric] for x in summ])
for metric in summ[0]
}
metrics_string = " ; ".join(
"{}: {:05.3f}".format(k, v) for k, v in metrics_mean.items())
logging.info("- [" + split.upper() + "] Eval metrics : " + metrics_string)
return metrics_mean