def train(dhandlers, enc, dec, d_hnet, device, config, writer):
""" Train replay model in continual fashion on MNIST dataset.
This is a helper function that loops over the range of tasks and
iteratively starts training the replay model on new tasks.
Args:
dhandlers: The dataset handlers.
enc: The model of the encoder network.
dec. The model of the decoder network.
d_hnet. The model of the decoder hyper network.
device: Torch device (cpu or gpu).
latent_sampler: An initialized distribution, we can sample from.
config: The command line arguments.
writer: The tensorboard summary writer.
"""
print('Training the MNIST replay model ...')
# get embedding lists for plotting
embd_list = init_plotting_embedding(dhandlers, d_hnet, writer, config)
# train the replay model task by task
for t in range(config.num_embeddings):
if config.replay_method == 'gan':
train_gan_one_t(dhandlers[t], enc, dec, d_hnet, device,
config, writer, embd_list, t)
else:
train_vae_one_t(dhandlers[t], enc, dec, d_hnet, device,
config, writer, embd_list, t)
def train_tasks(dhandlers_class, dhandlers_rp, enc, dec, d_hnet, class_net,
device, config, writer, infer_net=None):
""" Train continual learning experiments on MNIST dataset.
This is a helper function that loops over the range of tasks and
iteratively starts training the classifier and the replay model
on new tasks. Additionally, we save the task performace just after
training which can later be compared to the performance after training
on all tasks.
Args:
dhandlers_class: The dataset handlers for classification.
dhandlers_rp: The dataset handlers from the replay.
enc: The model of the encoder network.
dec: The model of the decoder network.
d_hnet. The model of the decoder hyper network.
class_net: The model of the classifier.
device: Torch device (cpu or gpu).
config: The command line arguments.
writer: The tensorboard summary writer.
infer_net: (optional) Task inference net, only used for testing.
Returns:
A list of test accuracies of all tasks directly after training.
"""
print('Training MNIST (task inference) classifier ...')
if not (config.upper_bound or (config.infer_task_id and
config.cl_scenario == 1)):
if not config.trained_replay_model:
embd_list = init_plotting_embedding(dhandlers_rp,
d_hnet, writer, config)
during_accs = []
# Begin training loop for the single tasks
for t in range(0, config.num_tasks):
dhandler = dhandlers_class[t]
if class_net is not None:
if not (config.class_incremental and t == 0):
print("Training classifier on data handler: ", t)
train_class_one_t(dhandler, dhandlers_rp, dec,
d_hnet, class_net, device, config, writer, t)
else:
if t > 0:
print("Training task inference system on data handler: ", t)
train_class_one_t(dhandler, dhandlers_rp, dec,
d_hnet, infer_net, device, config, writer, t)
if not (t == 0 and class_net is None):
durring_cc = test([dhandler], class_net, infer_net, device,
config, writer, task_id=t)
during_accs.append(durring_cc)
if not (config.upper_bound or (config.infer_task_id and
config.cl_scenario == 1)):
if not config.trained_replay_model and t < config.num_tasks - 1:
if config.replay_method == 'gan':
train_gan_one_t(dhandlers_rp[t], enc, dec, d_hnet, device,
config, writer, embd_list, t)
else:
train_vae_one_t(dhandlers_rp[t], enc, dec, d_hnet, device,
config, writer, embd_list, t)
return during_accs