def test_with_maml(dataset, learner, checkpoint, steps, loss_fn):
print("[*] Testing...")
model = MAML(learner, steps=steps, loss_function=loss_fn)
model.to(device)
if checkpoint:
model.restore(checkpoint, resume_training=False)
else:
print("[!] You are running inference on a randomly initialized model!")
model.eval(dataset, compute_accuracy=(type(dataset) is OmniglotDataset))
print("[*] Done!")
def train_with_maml(dataset,
learner,
save_path: str,
steps: int,
meta_batch_size: int,
iterations: int,
checkpoint=None,
loss_fn=None):
print("[*] Training...")
model = MAML(learner, steps=steps, loss_function=loss_fn)
model.to(device)
epoch = 0
if checkpoint:
model.restore(checkpoint)
epoch = checkpoint['epoch']
model.fit(dataset, iterations, save_path, epoch, 100)
print("[*] Done!")
return model
