def next(self, n_tasks, mode="train"):
if mode == 'train':
train_dataloader = BatchMetaDataLoader(self.train_dataset, batch_size=n_tasks, num_workers=0)
dataiter = iter(train_dataloader)
data = dataiter.next()
x_spts, y_spts = data["train"]
x_qrys, y_qrys = data["test"]
data = [x_spts.to(self.device), y_spts.to(self.device), x_qrys.to(self.device), y_qrys.to(self.device)]
return data
if mode == 'val':
val_dataloader = BatchMetaDataLoader(self.val_dataset, batch_size=n_tasks, num_workers=0)
dataiter = iter(val_dataloader)
data = dataiter.next()
x_spts, y_spts = data["train"]
x_qrys, y_qrys = data["test"]
data = [x_spts.to(self.device), y_spts.to(self.device), x_qrys.to(self.device), y_qrys.to(self.device)]
return data
if mode == 'test':
test_dataloader = BatchMetaDataLoader(self.test_dataset, batch_size=n_tasks, num_workers=0)
dataiter = iter(test_dataloader)
data = dataiter.next()
x_spts, y_spts = data["train"]
x_qrys, y_qrys = data["test"]
data = [x_spts.to(self.device), y_spts.to(self.device), x_qrys.to(self.device), y_qrys.to(self.device)]
return data
def get_torchmeta_rand_fnn_dataloaders(args):
# get data
dataset_train = RandFNN(args.data_path, 'train')
dataset_val = RandFNN(args.data_path, 'val')
dataset_test = RandFNN(args.data_path, 'test')
# get meta-sets
metaset_train = ClassSplitter(dataset_train,
num_train_per_class=args.k_shots,
num_test_per_class=args.k_eval,
shuffle=True)
metaset_val = ClassSplitter(dataset_val,
num_train_per_class=args.k_shots,
num_test_per_class=args.k_eval,
shuffle=True)
metaset_test = ClassSplitter(dataset_test,
num_train_per_class=args.k_shots,
num_test_per_class=args.k_eval,
shuffle=True)
# get meta-dataloader
meta_train_dataloader = BatchMetaDataLoader(
metaset_train,
batch_size=args.meta_batch_size_train,
num_workers=args.num_workers)
meta_val_dataloader = BatchMetaDataLoader(
metaset_val,
batch_size=args.meta_batch_size_eval,
num_workers=args.num_workers)
meta_test_dataloader = BatchMetaDataLoader(
metaset_test,
batch_size=args.meta_batch_size_eval,
num_workers=args.num_workers)
return meta_train_dataloader, meta_val_dataloader, meta_test_dataloader
def get_miniimagenet_dataloaders_torchmeta(args):
args.trainin_with_epochs = False
args.data_path = Path(
'~/data/').expanduser() # for some datasets this is enough
args.criterion = nn.CrossEntropyLoss()
# args.image_size = 84 # do we need this?
from torchmeta.datasets.helpers import miniimagenet
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
data_augmentation_transforms = transforms.Compose([
transforms.RandomResizedCrop(84),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.2),
transforms.ToTensor(), normalize
])
dataset_train = miniimagenet(args.data_path,
transform=data_augmentation_transforms,
ways=args.n_classes,
shots=args.k_shots,
test_shots=args.k_eval,
meta_split='train',
download=True)
dataset_val = miniimagenet(args.data_path,
ways=args.n_classes,
shots=args.k_shots,
test_shots=args.k_eval,
meta_split='val',
download=True)
dataset_test = miniimagenet(args.data_path,
ways=args.n_classes,
shots=args.k_shots,
test_shots=args.k_eval,
meta_split='test',
download=True)
meta_train_dataloader = BatchMetaDataLoader(
dataset_train,
batch_size=args.meta_batch_size_train,
num_workers=args.num_workers)
meta_val_dataloader = BatchMetaDataLoader(
dataset_val,
batch_size=args.meta_batch_size_eval,
num_workers=args.num_workers)
meta_test_dataloader = BatchMetaDataLoader(
dataset_test,
batch_size=args.meta_batch_size_eval,
num_workers=args.num_workers)
return meta_train_dataloader, meta_val_dataloader, meta_test_dataloader
def __init__(self,
dataset_name,
folder,
num_shots=5,
num_ways=5,
test_shots=15,
batch_size=20,
download=False,
dataloader_shuffle=True,
num_workers=0,
**kwargs):
datasetnamedict = {
'omniglot': omniglot,
'miniimagenet': miniimagenet,
'tieredimagenet': tieredimagenet
}
dataset = datasetnamedict[dataset_name]
#输入[metatrain,metatest]
test_shots = [test_shots, test_shots] if isinstance(
test_shots, int) else test_shots
batch_size = [batch_size, batch_size] if isinstance(
batch_size, int) else batch_size
#meta-train
self.train_dataset = dataset(folder,
shots=num_shots,
ways=num_ways,
shuffle=dataloader_shuffle,
test_shots=test_shots[0],
meta_train=True,
download=download)
self.train_dataloader = BatchMetaDataLoader(self.train_dataset,
batch_size=batch_size[0],
shuffle=dataloader_shuffle,
num_workers=num_workers)
#meta-test
self.test_dataset = dataset(folder,
shots=num_shots,
ways=num_ways,
shuffle=dataloader_shuffle,
test_shots=test_shots[1],
meta_test=True,
download=download)
self.test_dataloader = BatchMetaDataLoader(self.test_dataset,
batch_size=batch_size[1],
shuffle=dataloader_shuffle,
num_workers=num_workers)
def load_meta_trainset():
'''
dataset = Omniglot("data",
# Number of ways
num_classes_per_task=5,
# Resize the images to 28x28 and converts them to PyTorch tensors (from Torchvision)
transform=Compose([Resize(28), ToTensor()]),
# Transform the labels to integers (e.g. ("Glagolitic/character01", "Sanskrit/character14", ...) to (0, 1, ...))
target_transform=Categorical(num_classes=5),
# Creates new virtual classes with rotated versions of the images (from Santoro et al., 2016)
class_augmentations=[Rotation([90, 180, 270])],
meta_train=True,
download=True)
'''
trainset = omniglot("data", ways=config.n, shots=config.k, test_shots=15, shuffle=False, meta_train=True, download=True)
trainloader = BatchMetaDataLoader(trainset, batch_size=config.batch_size, shuffle=False, num_workers=0)
#trainset = Pascal5i("data", num_classes_per_task=config.n, meta_train=True, download=True)
#trainloader = BatchMetaDataLoader(dataset, batch_size=config.batch_size, shuffle=True, num_workers=0)
#trainset = CIFARFS("data", ways=config.n, shots=config.k, test_shots=15, shuffle=False, meta_train=True, download=True)
#trainloader = BatchMetaDataLoader(trainset, batch_size=config.batch_size, shuffle=False, num_workers=0)
return trainset, trainloader
def test_datasets_helpers_dataloader(name, shots, split):
function = getattr(helpers, name)
folder = os.getenv('TORCHMETA_DATA_FOLDER')
download = bool(os.getenv('TORCHMETA_DOWNLOAD', False))
dataset = function(folder,
ways=5,
shots=shots,
test_shots=15,
meta_split=split,
download=download)
meta_dataloader = BatchMetaDataLoader(dataset, batch_size=4)
batch = next(iter(meta_dataloader))
assert isinstance(batch, dict)
assert 'train' in batch
assert 'test' in batch
train_inputs, train_targets = batch['train']
test_inputs, test_targets = batch['test']
assert isinstance(train_inputs, torch.Tensor)
assert isinstance(train_targets, torch.Tensor)
assert train_inputs.ndim == 5
assert train_inputs.shape[:2] == (4, 5 * shots)
assert train_targets.ndim == 2
assert train_targets.shape[:2] == (4, 5 * shots)
assert isinstance(test_inputs, torch.Tensor)
assert isinstance(test_targets, torch.Tensor)
assert test_inputs.ndim == 5
assert test_inputs.shape[:2] == (4, 5 * 15) # test_shots
assert test_targets.ndim == 2
assert test_targets.shape[:2] == (4, 5 * 15)
def train(args):
logger.warning('This script is an example to showcase the extensions and '
'data-loading features of Torchmeta, and as such has been '
'very lightly tested.')
dataset = omniglot(args.folder,
shots=args.num_shots,
ways=args.num_ways,
shuffle=True,
test_shots=15,
meta_train=True,
download=args.download)
dataloader = BatchMetaDataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
model = PrototypicalNetwork(1,
args.embedding_size,
hidden_size=args.hidden_size)
model.to(device=args.device)
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# Training loop
with tqdm(dataloader, total=args.num_batches) as pbar:
for batch_idx, batch in enumerate(pbar):
model.zero_grad()
train_inputs, train_targets = batch['train']
train_inputs = train_inputs.to(device=args.device)
train_targets = train_targets.to(device=args.device)
train_embeddings = model(train_inputs)
test_inputs, test_targets = batch['test']
test_inputs = test_inputs.to(device=args.device)
test_targets = test_targets.to(device=args.device)
test_embeddings = model(test_inputs)
prototypes = get_prototypes(train_embeddings, train_targets,
dataset.num_classes_per_task)
loss = prototypical_loss(prototypes, test_embeddings, test_targets)
loss.backward()
optimizer.step()
with torch.no_grad():
accuracy = get_accuracy(prototypes, test_embeddings, test_targets)
pbar.set_postfix(accuracy='{0:.4f}'.format(accuracy.item()))
if batch_idx >= args.num_batches:
break
# Save model
if args.output_folder is not None:
filename = os.path.join(args.output_folder, 'protonet_omniglot_'
'{0}shot_{1}way.pt'.format(args.num_shots, args.num_ways))
with open(filename, 'wb') as f:
state_dict = model.state_dict()
torch.save(state_dict, f)
def get_torchmeta_sinusoid_dataloaders(args):
# tran = transforms.Compose([torch.tensor])
# dataset = sinusoid(shots=args.k_eval, test_shots=args.k_shots, transform=tran)
dataset = sinusoid(shots=args.k_eval, test_shots=args.k_eval)
meta_train_dataloader = BatchMetaDataLoader(
dataset,
batch_size=args.meta_batch_size_train,
num_workers=args.num_workers)
meta_val_dataloader = BatchMetaDataLoader(
dataset,
batch_size=args.meta_batch_size_eval,
num_workers=args.num_workers)
meta_test_dataloader = BatchMetaDataLoader(
dataset,
batch_size=args.meta_batch_size_eval,
num_workers=args.num_workers)
return meta_train_dataloader, meta_val_dataloader, meta_test_dataloader
def prepare_dataset(args, transform):
dataset = get_fewshotsen12msdataset(args.dataset_path, shots=args.num_shots, ways=args.num_ways, transform=transform,
target_transform=None,
meta_split="train", shuffle=True)
dataloader = BatchMetaDataLoader(dataset, batch_size=args.batch_size,
shuffle=False, num_workers=args.num_workers,
sampler=CombinationSubsetRandomSampler(dataset))
valdataset = get_fewshotsen12msdataset(args.dataset_path, shots=args.num_shots, ways=args.num_ways, transform=transform,
target_transform=None,
meta_split="val", shuffle=True)
valdataloader = BatchMetaDataLoader(valdataset, batch_size=args.batch_size,
shuffle=False, num_workers=args.num_workers,
sampler=CombinationSubsetRandomSampler(valdataset))
return dataloader, valdataloader
def train(args):
logger.warning(
'This script is an example to showcase the MetaModule and '
'data-loading features of Torchmeta, and as such has been '
'very lightly tested. For a better tested implementation of '
'Model-Agnostic Meta-Learning (MAML) using Torchmeta with '
'more features (including multi-step adaptation and '
'different datasets), please check `https://github.com/'
'tristandeleu/pytorch-maml`.')
inputs = [
"new_models/210429/dataset/BedBathingBaxterHuman-v0217_0-v1-human-coop-robot-coop_10k"
]
env_name = "BedBathingBaxterHuman-v0217_0-v1"
env = gym.make('assistive_gym:' + env_name)
dataset = behaviour(inputs, shots=400, test_shots=1)
dataloader = BatchMetaDataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
model = PolicyNetwork(env.observation_space_human.shape[0],
env.action_space_human.shape[0])
for key, v in model.features.named_parameters():
v.data = torch.nn.init.zeros_(v)
model.to(device=args.device)
model.train()
meta_optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# Training loop
pbar = tqdm(total=args.num_batches)
batch_idx = 0
while batch_idx < args.num_batches:
for batch in dataloader:
model.zero_grad()
train_inputs, train_targets = batch['train']
train_inputs = train_inputs.to(device=args.device).float()
train_targets = train_targets.to(device=args.device).float()
loss = torch.tensor(0., device=args.device)
for task_idx, (train_input, train_target) in enumerate(
zip(train_inputs, train_targets)):
train_output = model(train_input)
loss += get_loss(train_output, train_target)
model.zero_grad()
loss.div_(len(dataloader))
loss.backward()
meta_optimizer.step()
pbar.update(1)
pbar.set_postfix(loss='{0:.4f}'.format(loss.item()))
batch_idx += 1
def main(
shots=10,
tasks_per_batch=16,
num_tasks=160000,
adapt_lr=0.01,
meta_lr=0.001,
adapt_steps=5,
hidden_dim=32,
):
# load the dataset
tasksets = Sinusoid(num_samples_per_task=2 * shots, num_tasks=num_tasks)
dataloader = BatchMetaDataLoader(tasksets, batch_size=tasks_per_batch)
# create the model
model = SineModel(dim=hidden_dim)
maml = l2l.algorithms.MAML(model,
lr=adapt_lr,
first_order=False,
allow_unused=True)
opt = optim.Adam(maml.parameters(), meta_lr)
lossfn = nn.MSELoss(reduction='mean')
# for each iteration
for iter, batch in enumerate(dataloader): # num_tasks/batch_size
meta_train_loss = 0.0
# for each task in the batch
effective_batch_size = batch[0].shape[0]
for i in range(effective_batch_size):
learner = maml.clone()
# divide the data into support and query sets
train_inputs, train_targets = batch[0][i].float(
), batch[1][i].float()
x_support, y_support = train_inputs[::2], train_targets[::2]
x_query, y_query = train_inputs[1::2], train_targets[1::2]
for _ in range(adapt_steps): # adaptation_steps
support_preds = learner(x_support)
support_loss = lossfn(support_preds, y_support)
learner.adapt(support_loss)
query_preds = learner(x_query)
query_loss = lossfn(query_preds, y_query)
meta_train_loss += query_loss
meta_train_loss = meta_train_loss / effective_batch_size
if iter % 200 == 0:
print('Iteration:', iter, 'Meta Train Loss',
meta_train_loss.item())
opt.zero_grad()
meta_train_loss.backward()
opt.step()
def prepare_fewshotdataloader(root, shots, ways, fold, transform, shuffle=True, num_tasks=1,
num_workers=0, download=False):
dataset = get_fewshotsen12msdataset(root, shots=shots, ways=ways, transform=transform,
target_transform=None,
meta_split=fold, shuffle=shuffle, download=download)
dataloader = BatchMetaDataLoader(dataset, batch_size=num_tasks,
shuffle=False, num_workers=num_workers,
sampler=CombinationSubsetRandomSampler(dataset))
return dataloader
def load_meta_testset():
testset = omniglot("data", ways=config.n, shots=config.k, test_shots=15, shuffle=False, meta_test=True, download=True)
testloader = BatchMetaDataLoader(testset, batch_size=config.batch_size, shuffle=False, num_workers=0)
#testset = Pascal5i("data", num_classes_per_task=config.n, meta_test=True, download=True)
#testloader = BatchMetaDataLoader(dataset, batch_size=config.batch_size, shuffle=True, num_workers=0)
#testset = CIFARFS("data", ways=config.n, shots=config.k, test_shots=15, shuffle=False, meta_test=True, download=True)
#testloader = BatchMetaDataLoader(testset, batch_size=config.batch_size, shuffle=False, num_workers=0)
return testset, testloader
def test_batch_meta_dataloader():
dataset = Sinusoid(10, num_tasks=1000, noise_std=None)
meta_dataloader = BatchMetaDataLoader(dataset, batch_size=4)
assert isinstance(meta_dataloader, DataLoader)
assert len(meta_dataloader) == 250 # 1000 / 4
inputs, targets = next(iter(meta_dataloader))
assert isinstance(inputs, torch.Tensor)
assert isinstance(targets, torch.Tensor)
assert inputs.shape == (4, 10, 1)
assert targets.shape == (4, 10, 1)
def create_og_data_loader(
root,
meta_split,
k_way,
n_shot,
input_size,
n_query,
batch_size,
num_workers,
download=False,
use_vinyals_split=False,
seed=None,
):
"""Create a torchmeta BatchMetaDataLoader for Omniglot
Args:
root: Path to Omniglot data root folder (containing an 'omniglot'` subfolder with the
preprocess json-Files or downloaded zip-files).
meta_split: see torchmeta.datasets.Omniglot
k_way: Number of classes per task
n_shot: Number of samples per class
input_size: Images are resized to this size.
n_query: Number of test images per class
batch_size: Meta batch size
num_workers: Number of workers for data preprocessing
download: Download (and dataset specific preprocessing that needs to be done on the
downloaded files).
use_vinyals_split: see torchmeta.datasets.Omniglot
seed: Seed to be used in the meta-dataset
Returns:
A torchmeta :class:`BatchMetaDataLoader` object.
"""
dataset = Omniglot(
root,
num_classes_per_task=k_way,
transform=Compose([Resize(input_size), ToTensor()]),
target_transform=Categorical(num_classes=k_way),
class_augmentations=[Rotation([90, 180, 270])],
meta_split=meta_split,
download=download,
use_vinyals_split=use_vinyals_split,
)
dataset = ClassSplitter(dataset,
shuffle=True,
num_train_per_class=n_shot,
num_test_per_class=n_query)
dataset.seed = seed
dataloader = BatchMetaDataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
return dataloader
def main(args):
test_dataset = dataset_f(args, 'test')
if args.support_samples > 0:
test_dloader = BatchMetaDataLoader(test_dataset, args.batch_size)
else:
test_dloader = DataLoader(test_dataset, args.batch_size)
model_checkpoint = load_best_model_checkpoint(args.run_path, args.device)
zero_shot = args.support_samples == 0
if zero_shot:
model = PrototypicalNetworkZeroShot(
args.distance,
num_classes=args.num_classes,
meta_features=args.train_config['metadata_features'],
img_features=args.train_config['image_features'])
else:
input_channels = 1 if args.dataset == 'omniglot' else 3
model = PrototypicalNetwork(args.distance,
args.num_classes,
input_channels=input_channels)
model.load_state_dict(model_checkpoint)
model = model.float().to(args.device)
evaluator = Trainer(model,
None,
test_dloader,
args.distance,
args.run_path,
train_epochs=-1,
opt=None,
device=args.device,
eval_steps=args.steps,
zero_shot=zero_shot)
test_results = evaluator.eval()['val']
print('evaluation done')
eval_folder: Path = args.run_path / 'evaluation'
eval_folder.mkdir_p()
results = dict(num_classes=args.num_classes,
support_samples=args.support_samples,
query_samples=args.query_samples,
batch_size=args.batch_size,
steps=args.steps,
eval_duration_seconds=test_results.eval_duration_seconds,
avg_accuracy=test_results.metrics['accuracy'],
avg_loss=test_results.metrics['avg_loss'])
dst_file = eval_folder / \
f'num_classes={args.num_classes}_' \
f'support={args.support_samples}_' \
f'query={args.query_samples}.json'
with open(dst_file, 'w') as f:
json.dump(results, f, indent=2)
print('saved results to', dst_file)
def main(args):
with open(args.config, 'r') as f:
config = json.load(f)
if args.folder is not None:
config['folder'] = args.folder
if args.num_steps > 0:
config['num_steps'] = args.num_steps
if args.num_batches > 0:
config['num_batches'] = args.num_batches
config_dir = path.dirname(args.config)
config['folder'] = path.join(config_dir, config['folder'])
config['model_path'] = path.join(config_dir, config['model_path'])
device = torch.device('cuda' if args.use_cuda and torch.cuda.is_available() else 'cpu')
benchmark = get_benchmark_by_name(config['dataset'],
config['folder'],
config['num_ways'],
config['num_shots'],
config['num_shots_test'],
config['no_max_pool'],
hidden_size=config['hidden_size'])
with open(config['model_path'], 'rb') as f:
benchmark.model.load_state_dict(torch.load(f, map_location=device))
meta_test_dataloader = BatchMetaDataLoader(benchmark.meta_test_dataset,
batch_size=config['batch_size'],
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
metalearner = ModelAgnosticMetaLearning(benchmark.model,
first_order=config['first_order'],
num_adaptation_steps=config['num_steps'],
step_size=config['step_size'],
loss_function=benchmark.loss_function,
device=device)
results = metalearner.evaluate(meta_test_dataloader,
max_batches=config['num_batches'],
silent=args.silent,
desc='Test')
# Save results
dirname = os.path.dirname(config['model_path'])
with open(os.path.join(dirname, 'results.json'), 'w') as f:
json.dump(results, f)
def get_dataset_loader(dataset_id,
folder,
shot,
query_size,
batch_size,
shuffle,
train=False,
val=False,
test=False):
dataset = get_dataset(dataset_id, folder, shot, query_size, shuffle, train,
val, test)
loader = BatchMetaDataLoader(dataset,
batch_size=batch_size,
shuffle=shuffle)
return loader
def get_sine_loader(batch_size, num_steps, shots=10, test_shots=15):
dataset_transform = ClassSplitter(
shuffle=True, num_train_per_class=shots, num_test_per_class=test_shots
)
transform = ToTensor1D()
dataset = Sinusoid(
shots + test_shots,
num_tasks=batch_size * num_steps,
transform=transform,
target_transform=transform,
dataset_transform=dataset_transform,
)
loader = BatchMetaDataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True,
)
return loader
def helper_meta_imagelist(self, batch_size, input_size, num_shots,
num_shots_test, num_ways, num_workers):
with tempfile.TemporaryDirectory(
) as folder, tempfile.NamedTemporaryFile(mode='w+t') as fp:
create_random_imagelist(folder, fp, input_size, create_image)
dataset = data.ImagelistMetaDataset(
imagelistname=fp.name,
root='',
transform=transforms.Compose(
[transforms.Resize(input_size),
transforms.ToTensor()]))
meta_dataset = CombinationMetaDataset(
dataset,
num_classes_per_task=num_ways,
target_transform=Categorical(num_ways),
dataset_transform=ClassSplitter(
shuffle=True,
num_train_per_class=num_shots,
num_test_per_class=num_shots_test))
meta_dataloader = BatchMetaDataLoader(meta_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True)
for batch in meta_dataloader:
batch_data, batch_label = batch['train']
for img_train, label_train, img_test, label_test in zip(
*batch['train'], *batch['test']):
classmap = {}
for idx in range(img_train.shape[0]):
npimg = img_train[idx, ...].detach().numpy()
npimg[npimg < 0.001] = 0
imagesum = int(np.sum(npimg) * 255)
if imagesum not in classmap:
classmap[imagesum] = int(label_train[idx])
self.assertEqual(classmap[imagesum],
int(label_train[idx]))
for idx in range(img_test.shape[0]):
npimg = img_test[idx, ...].detach().numpy()
npimg[npimg < 0.001] = 0
imagesum = int(np.sum(npimg) * 255)
self.assertEqual(classmap[imagesum],
int(label_test[idx]),
"Error on {}".format(idx))
def test_overflow_length_dataloader():
folder = os.getenv('TORCHMETA_DATA_FOLDER')
download = bool(os.getenv('TORCHMETA_DOWNLOAD', False))
# The number of tasks is C(4112, 20), which exceeds machine precision
dataset = helpers.omniglot(folder,
ways=20,
shots=1,
test_shots=5,
meta_train=True,
download=download)
meta_dataloader = BatchMetaDataLoader(dataset, batch_size=4)
batch = next(iter(meta_dataloader))
assert isinstance(batch, dict)
assert 'train' in batch
assert 'test' in batch
def create_miniimagenet_data_loader(
root,
meta_split,
k_way,
n_shot,
n_query,
batch_size,
num_workers,
download=False,
seed=None,
):
"""Create a torchmeta BatchMetaDataLoader for MiniImagenet
Args:
root: Path to mini imagenet root folder (containing an 'miniimagenet'` subfolder with the
preprocess json-Files or downloaded tar.gz-file).
meta_split: see torchmeta.datasets.MiniImagenet
k_way: Number of classes per task
n_shot: Number of samples per class
n_query: Number of test images per class
batch_size: Meta batch size
num_workers: Number of workers for data preprocessing
download: Download (and dataset specific preprocessing that needs to be done on the
downloaded files).
seed: Seed to be used in the meta-dataset
Returns:
A torchmeta :class:`BatchMetaDataLoader` object.
"""
dataset = miniimagenet(
root,
n_shot,
k_way,
meta_split=meta_split,
test_shots=n_query,
download=download,
seed=seed,
)
dataloader = BatchMetaDataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
return dataloader
def test_batch_meta_dataloader_splitter():
dataset = Sinusoid(20, num_tasks=1000, noise_std=None)
dataset = ClassSplitter(dataset, num_train_per_class=5,
num_test_per_class=15)
meta_dataloader = BatchMetaDataLoader(dataset, batch_size=4)
batch = next(iter(meta_dataloader))
assert isinstance(batch, dict)
assert 'train' in batch
assert 'test' in batch
train_inputs, train_targets = batch['train']
test_inputs, test_targets = batch['test']
assert isinstance(train_inputs, torch.Tensor)
assert isinstance(train_targets, torch.Tensor)
assert train_inputs.shape == (4, 5, 1)
assert train_targets.shape == (4, 5, 1)
assert isinstance(test_inputs, torch.Tensor)
assert isinstance(test_targets, torch.Tensor)
assert test_inputs.shape == (4, 15, 1)
assert test_targets.shape == (4, 15, 1)
# meta_test_ds = Omniglot(root_path, num_classes_per_task=5, meta_test=True,
# use_vinyals_split=False, download=True, transform=trans)
return (meta_train_ds, meta_test_ds)
if __name__ == '__main__':
# (meta_train_ds, meta_test_ds) = get_omniglot_for_fl()
from torchmeta.datasets.helpers import omniglot
from torchmeta.utils.data import BatchMetaDataLoader
dataset = omniglot(root_path,
ways=5,
shots=5,
test_shots=15,
meta_train=True,
download=True)
dataloader = BatchMetaDataLoader(dataset, batch_size=16, num_workers=4)
for batch in dataloader:
train_inputs, train_targets = batch["train"]
print('Train inputs shape: {0}'.format(
train_inputs.shape)) # (16, 25, 1, 28, 28)
print('Train targets shape: {0}'.format(
train_targets.shape)) # (16, 25)
test_inputs, test_targets = batch["test"]
print('Test inputs shape: {0}'.format(
test_inputs.shape)) # (16, 75, 1, 28, 28)
print('Test targets shape: {0}'.format(test_targets.shape)) # (16, 75)
g = 5
def main(args, mode, iteration=None):
dataset = load_dataset(args, mode)
dataloader = BatchMetaDataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
model.to(device=args.device)
model.train()
# To control outer update parameter
# If you want to control inner update parameter, please see update_parameters function in ./maml/utils.py
freeze_params = [
p for name, p in model.named_parameters() if 'classifier' in name
]
learnable_params = [
p for name, p in model.named_parameters() if 'classifier' not in name
]
if args.outer_fix:
meta_optimizer = torch.optim.Adam([{
'params': freeze_params,
'lr': 0
}, {
'params': learnable_params,
'lr': args.meta_lr
}])
else:
meta_optimizer = torch.optim.Adam([{
'params': freeze_params,
'lr': args.meta_lr
}, {
'params': learnable_params,
'lr': args.meta_lr
}])
if args.meta_train:
total = args.train_batches
elif args.meta_val:
total = args.valid_batches
elif args.meta_test:
total = args.test_batches
loss_logs, accuracy_logs = [], []
# Training loop
with tqdm(dataloader, total=total, leave=False) as pbar:
for batch_idx, batch in enumerate(pbar):
if args.centering:
fc_weight_mean = torch.mean(model.classifier.weight.data,
dim=0)
model.classifier.weight.data -= fc_weight_mean
model.zero_grad()
support_inputs, support_targets = batch['train']
support_inputs = support_inputs.to(device=args.device)
support_targets = support_targets.to(device=args.device)
query_inputs, query_targets = batch['test']
query_inputs = query_inputs.to(device=args.device)
query_targets = query_targets.to(device=args.device)
outer_loss = torch.tensor(0., device=args.device)
accuracy = torch.tensor(0., device=args.device)
for task_idx, (support_input, support_target, query_input,
query_target) in enumerate(
zip(support_inputs, support_targets,
query_inputs, query_targets)):
support_features, support_logit = model(support_input)
inner_loss = F.cross_entropy(support_logit, support_target)
model.zero_grad()
params = update_parameters(
model,
inner_loss,
extractor_step_size=args.extractor_step_size,
classifier_step_size=args.classifier_step_size,
first_order=args.first_order)
query_features, query_logit = model(query_input, params=params)
outer_loss += F.cross_entropy(query_logit, query_target)
with torch.no_grad():
accuracy += get_accuracy(query_logit, query_target)
outer_loss.div_(args.batch_size)
accuracy.div_(args.batch_size)
loss_logs.append(outer_loss.item())
accuracy_logs.append(accuracy.item())
if args.meta_train:
outer_loss.backward()
meta_optimizer.step()
postfix = {
'mode': mode,
'iter': iteration,
'acc': round(accuracy.item(), 5)
}
pbar.set_postfix(postfix)
if batch_idx + 1 == total:
break
# Save model
if args.meta_train:
filename = os.path.join(args.output_folder,
args.dataset + '_' + args.save_name, 'models',
'epochs_{}.pt'.format((iteration + 1) * total))
if (iteration + 1) * total % 5000 == 0:
with open(filename, 'wb') as f:
state_dict = model.state_dict()
torch.save(state_dict, f)
return loss_logs, accuracy_logs
def main(args):
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
device = torch.device(
'cuda' if args.use_cuda and torch.cuda.is_available() else 'cpu')
if (args.output_folder is not None):
if not os.path.exists(args.output_folder):
os.makedirs(args.output_folder)
logging.debug('Creating folder `{0}`'.format(args.output_folder))
folder = os.path.join(args.output_folder,
time.strftime('%Y-%m-%d_%H%M%S'))
os.makedirs(folder)
logging.debug('Creating folder `{0}`'.format(folder))
args.folder = os.path.abspath(args.folder)
args.model_path = os.path.abspath(os.path.join(folder, 'model.th'))
# Save the configuration in a config.json file
with open(os.path.join(folder, 'config.json'), 'w') as f:
json.dump(vars(args), f, indent=2)
logging.info('Saving configuration file in `{0}`'.format(
os.path.abspath(os.path.join(folder, 'config.json'))))
dataset_transform = ClassSplitter(shuffle=True,
num_train_per_class=args.num_shots,
num_test_per_class=args.num_shots_test)
class_augmentations = [Rotation([90, 180, 270])]
if args.dataset == 'sinusoid':
transform = ToTensor()
meta_train_dataset = Sinusoid(args.num_shots + args.num_shots_test,
num_tasks=1000000,
transform=transform,
target_transform=transform,
dataset_transform=dataset_transform)
meta_val_dataset = Sinusoid(args.num_shots + args.num_shots_test,
num_tasks=1000000,
transform=transform,
target_transform=transform,
dataset_transform=dataset_transform)
model = ModelMLPSinusoid(hidden_sizes=[40, 40])
loss_function = F.mse_loss
elif args.dataset == 'omniglot':
transform = Compose([Resize(28), ToTensor()])
meta_train_dataset = Omniglot(args.folder,
transform=transform,
target_transform=Categorical(
args.num_ways),
num_classes_per_task=args.num_ways,
meta_train=True,
class_augmentations=class_augmentations,
dataset_transform=dataset_transform,
download=True)
meta_val_dataset = Omniglot(args.folder,
transform=transform,
target_transform=Categorical(
args.num_ways),
num_classes_per_task=args.num_ways,
meta_val=True,
class_augmentations=class_augmentations,
dataset_transform=dataset_transform)
model = ModelConvOmniglot(args.num_ways, hidden_size=args.hidden_size)
loss_function = F.cross_entropy
elif args.dataset == 'miniimagenet':
transform = Compose([Resize(84), ToTensor()])
meta_train_dataset = MiniImagenet(
args.folder,
transform=transform,
target_transform=Categorical(args.num_ways),
num_classes_per_task=args.num_ways,
meta_train=True,
class_augmentations=class_augmentations,
dataset_transform=dataset_transform,
download=True)
meta_val_dataset = MiniImagenet(
args.folder,
transform=transform,
target_transform=Categorical(args.num_ways),
num_classes_per_task=args.num_ways,
meta_val=True,
class_augmentations=class_augmentations,
dataset_transform=dataset_transform)
model = ModelConvMiniImagenet(args.num_ways,
hidden_size=args.hidden_size)
loss_function = F.cross_entropy
else:
raise NotImplementedError('Unknown dataset `{0}`.'.format(
args.dataset))
meta_train_dataloader = BatchMetaDataLoader(meta_train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
meta_val_dataloader = BatchMetaDataLoader(meta_val_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
meta_optimizer = torch.optim.Adam(model.parameters(), lr=args.meta_lr)
metalearner = ModelAgnosticMetaLearning(
model,
meta_optimizer,
first_order=args.first_order,
num_adaptation_steps=args.num_steps,
step_size=args.step_size,
loss_function=loss_function,
device=device)
best_val_accuracy = None
# Training loop
epoch_desc = 'Epoch {{0: <{0}d}}'.format(1 +
int(math.log10(args.num_epochs)))
for epoch in range(args.num_epochs):
metalearner.train(meta_train_dataloader,
max_batches=args.num_batches,
verbose=args.verbose,
desc='Training',
leave=False)
results = metalearner.evaluate(meta_val_dataloader,
max_batches=args.num_batches,
verbose=args.verbose,
desc=epoch_desc.format(epoch + 1))
if (best_val_accuracy is None) \
or (best_val_accuracy < results['accuracies_after']):
best_val_accuracy = results['accuracies_after']
if args.output_folder is not None:
with open(args.model_path, 'wb') as f:
torch.save(model.state_dict(), f)
if hasattr(meta_train_dataset, 'close'):
meta_train_dataset.close()
meta_val_dataset.close()
def main(
shots=10,
tasks_per_batch=16,
num_tasks=16000,
num_test_tasks=32,
adapt_lr=0.01,
meta_lr=0.001,
adapt_steps=5,
hidden_dim=32,
):
exp_name = input(
"Enter Experiment NAME (should be unique, else overwrites): ")
EXPERIMENT_DIR = "./experiments/MAML_Sine_exps" + exp_name
if os.path.isdir(EXPERIMENT_DIR):
print("Experiment folder opened ...")
else:
os.mkdir(EXPERIMENT_DIR)
print("New Experiment folder started ...")
MODEL_CHECKPOINT_DIR = EXPERIMENT_DIR + "/model"
if os.path.isdir(MODEL_CHECKPOINT_DIR):
print("Model Checkpoint folder opened ...")
else:
os.mkdir(MODEL_CHECKPOINT_DIR)
print("New Model checkpoint folder made ...")
PLOT_RESULTS_DIR = EXPERIMENT_DIR + "/plot_results"
if os.path.isdir(PLOT_RESULTS_DIR):
print("Image results folder opened ...")
else:
os.mkdir(PLOT_RESULTS_DIR)
print("New Image results folder made ...")
# load the dataset
tasksets = Sinusoid(num_samples_per_task=2 * shots, num_tasks=num_tasks)
dataloader = BatchMetaDataLoader(tasksets, batch_size=tasks_per_batch)
# create the model
model = SineModel(dim=hidden_dim, experiment_dir=EXPERIMENT_DIR)
maml = l2l.algorithms.MAML(model,
lr=adapt_lr,
first_order=False,
allow_unused=True)
opt = optim.Adam(maml.parameters(), meta_lr)
lossfn = nn.MSELoss(reduction='mean')
# for each iteration
for iter, batch in enumerate(dataloader): # num_tasks/batch_size
meta_train_loss = 0.0
# for each task in the batch
effective_batch_size = batch[0].shape[0]
for i in range(effective_batch_size):
learner = maml.clone()
# divide the data into support and query sets
train_inputs, train_targets = batch[0][i].float(
), batch[1][i].float()
x_support, y_support = train_inputs[::2], train_targets[::2]
x_query, y_query = train_inputs[1::2], train_targets[1::2]
for _ in range(adapt_steps): # adaptation_steps
support_preds = learner(x_support)
support_loss = lossfn(support_preds, y_support)
learner.adapt(support_loss)
query_preds = learner(x_query)
query_loss = lossfn(query_preds, y_query)
meta_train_loss += query_loss
meta_train_loss = meta_train_loss / effective_batch_size
opt.zero_grad()
meta_train_loss.backward()
opt.step()
if iter % 100 == 0:
print('Iteration:', iter, 'Meta Train Loss',
meta_train_loss.item())
# print(x_query.requires_grad, y_query.requires_grad,query_preds.requires_grad,meta_train_loss.item())
plotter(x_query, y_query,
query_preds.detach().numpy(), iter, 'Train',
meta_train_loss.item(), model.plot_results)
#save current model
model.save_checkpoint()
#meta-testing
test_tasks = Sinusoid(num_samples_per_task=shots, num_tasks=num_test_tasks)
test_dataloader = BatchMetaDataLoader(test_tasks,
batch_size=tasks_per_batch)
#load learned model
test_model = SineModel(dim=hidden_dim, experiment_dir=EXPERIMENT_DIR)
test_model.load_checkpoint()
for iter, batch in enumerate(test_dataloader):
meta_test_loss = 0.0
# for each task in the batch
effective_batch_size = batch[0].shape[0]
for i in range(effective_batch_size):
learner = maml.clone()
# divide the data into support and query sets
test_inputs, test_targets = batch[0][i].float(), batch[1][i].float(
)
test_preds = test_model(test_inputs)
test_loss = lossfn(test_preds, test_targets)
meta_test_loss += test_loss
meta_test_loss = meta_test_loss / effective_batch_size
if iter % 20 == 0:
print('Iteration:', iter, 'Meta Test Loss', meta_test_loss.item())
plotter(test_inputs, test_targets,
test_preds.detach().numpy(), iter, 'Test',
meta_test_loss.item(), test_model.plot_results)
def do_train(self, train_data, test_data, batch_size=0, accumulation_steps=1, num_epochs=1, test_strategy='all',
test_batch_size=0, include_query_point=True, mixup_alpha=None, verbose=True):
"""
Performs `num_epochs` training iterations of meta-NML on a batch of data.
Parameters
----------
train_data : tuple[np.array, np.array]
Inputs and labels to sample from for adaptation (inner loop).
Each task will consist of one of these inputs with a proposed class label.
test_data : tuple[np.array, np.array]
Inputs and labels to use for the test loss (outer loop).
The loss will be computed over ALL of these points for each task.
batch_size : int
Number of tasks to use per batch for meta-learning
num_epochs : int
Number of passes through the entire `train_data` dataset
test_strategy : str in {'all', 'sample', 'cycle'}
The strategy to use for evaluating the test loss across tasks in a batch.
By default, we use all the points in `test_data`.
test_batch_size : int, optional
Number of points to use in a test batch.
Only used if test_strategy is 'sample' or 'cycle'.
include_query_point : bool, optional
Whether to include the downweighted query point in every batch during testing. Only used
if the test_strategy is 'sample' or 'cycle'. Default: True
mixup_alpha : float, optional
Alpha parameter to use for mixup (https://arxiv.org/pdf/1710.09412.pdf)
(only affects the test set, not the tasks themselves).
verbose : bool
Whether to show the MAML training progress bar for each epoch. Default: True
"""
batch_size = batch_size or len(train_data[0])
self.model.train()
if not self._do_metalearning:
# Do standard MLE training on the meta-test set
ds = TensorDataset(torch.Tensor(test_data[0]), torch.Tensor(test_data[1]).long())
loader = DataLoader(ds, batch_size=64, shuffle=True)
epoch_results = []
for _ in range(num_epochs):
all_losses = []
for inputs, labels in loader:
inputs, labels = inputs.to(self.device), labels.to(self.device)
logits = self.model(inputs)
loss = F.cross_entropy(logits, labels)
self.meta_optimizer.zero_grad()
loss.backward()
self.meta_optimizer.step()
all_losses.append(loss.item())
epoch_results.append({
'mean_loss': np.array(all_losses).mean(),
'all_losses': all_losses
})
return epoch_results
epoch_results = []
for _ in range(num_epochs):
if self.embedding_type == 'features':
train_features = self.model.embedding(torch.Tensor(train_data[0]).cuda()).cpu().detach()
test_features = self.model.embedding(torch.Tensor(test_data[0]).cuda()).cpu().detach()
elif self.embedding_type == 'vae':
train_features = self.model_vae(torch.Tensor(train_data[0]).cuda())[1].cpu().detach()
test_features = self.model_vae(torch.Tensor(test_data[0]).cuda())[1].cpu().detach()
elif self.embedding_type == 'custom':
train_features = train_data[2]
test_features = test_data[2]
else:
train_features = train_data[0]
test_features = test_data[0]
train_data = (train_data[0], train_data[1], train_features)
test_data = (test_data[0], test_data[1], test_features)
dataset = NMLDataset(train_data, test_data, mixup_alpha=mixup_alpha,
points_per_task=self.points_per_task, num_classes=self.num_classes,
test_strategy=test_strategy,
test_batch_size=test_batch_size, include_query_point=include_query_point,
dist_weight_thresh=self.dist_weight_thresh, equal_pos_neg_test=self.equal_pos_neg_test,
query_point_weight=self.query_point_weight, kernel=self.kernel)
trainloader = BatchMetaDataLoader(dataset,
batch_size=batch_size, shuffle=True,
num_workers=self.num_workers, pin_memory=True)
num_batches = ceil(dataset.num_tasks / batch_size)
results = self.metalearner.train(trainloader, accumulation_steps=accumulation_steps, max_batches=num_batches, is_classification_task=True,
verbose=verbose, desc='Training', leave=False)
epoch_results.append(results)
return epoch_results
global_task_count = 0
# save the args into .json file
with open(os.path.join(args.record_folder, 'args.json'), 'w') as f:
json.dump(vars(args), f)
# get datasets and dataloaders
train_dataset = get_dataset(args,
dataset_name=args.train_data,
phase='train')
val_dataset = get_dataset(args, dataset_name=args.test_data, phase='val')
test_dataset = get_dataset(args, dataset_name=args.test_data, phase='test')
train_loader = BatchMetaDataLoader(train_dataset,
batch_size=args.batch_tasks,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
val_loader = BatchMetaDataLoader(val_dataset,
batch_size=args.batch_tasks,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True)
test_loader = BatchMetaDataLoader(test_dataset,
batch_size=args.batch_tasks,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True)
def main(args):
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
device = torch.device(
'cuda' if args.use_cuda and torch.cuda.is_available() else 'cpu')
# Create output folder
if (args.output_folder is not None):
if not os.path.exists(args.output_folder):
os.makedirs(args.output_folder)
logging.debug('Creating folder `{0}`'.format(args.output_folder))
output_folder = os.path.join(args.output_folder,
time.strftime('%Y-%m-%d_%H%M%S'))
os.makedirs(output_folder)
logging.debug('Creating folder `{0}`'.format(output_folder))
args.datafolder = os.path.abspath(args.datafolder)
args.model_path = os.path.abspath(
os.path.join(output_folder, 'model.th'))
# Save the configuration in a config.json file
with open(os.path.join(output_folder, 'config.json'), 'w') as f:
json.dump(vars(args), f, indent=2)
logging.info('Saving configuration file in `{0}`'.format(
os.path.abspath(os.path.join(output_folder, 'config.json'))))
# Get datasets and load into meta learning format
meta_train_dataset, meta_val_dataset, _ = get_datasets(
args.dataset,
args.datafolder,
args.num_ways,
args.num_shots,
args.num_shots_test,
augment=augment,
fold=args.fold,
download=download_data)
meta_train_dataloader = BatchMetaDataLoader(meta_train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
meta_val_dataloader = BatchMetaDataLoader(meta_val_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
# Define model
model = Unet(device=device, feature_scale=args.feature_scale)
model = model.to(device)
print(f'Using device: {device}')
# Define optimizer
meta_optimizer = torch.optim.Adam(model.parameters(),
lr=args.meta_lr) #, weight_decay=1e-5)
#meta_optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate, momentum = 0.99)
# Define meta learner
metalearner = ModelAgnosticMetaLearning(
model,
meta_optimizer,
first_order=args.first_order,
num_adaptation_steps=args.num_adaption_steps,
step_size=args.step_size,
learn_step_size=False,
loss_function=loss_function,
device=device)
best_value = None
# Training loop
epoch_desc = 'Epoch {{0: <{0}d}}'.format(1 +
int(math.log10(args.num_epochs)))
train_losses = []
val_losses = []
train_ious = []
train_accuracies = []
val_accuracies = []
val_ious = []
start_time = time.time()
for epoch in range(args.num_epochs):
print('start epoch ', epoch + 1)
print('start train---------------------------------------------------')
train_loss, train_accuracy, train_iou = metalearner.train(
meta_train_dataloader,
max_batches=args.num_batches,
verbose=args.verbose,
desc='Training',
leave=False)
print(f'\n train accuracy: {train_accuracy}, train loss: {train_loss}')
print('end train---------------------------------------------------')
train_losses.append(train_loss)
train_accuracies.append(train_accuracy)
train_ious.append(train_iou)
# Evaluate in given intervals
if epoch % args.val_step_size == 0:
print(
'start evaluate-------------------------------------------------'
)
results = metalearner.evaluate(meta_val_dataloader,
max_batches=args.num_batches,
verbose=args.verbose,
desc=epoch_desc.format(epoch + 1),
is_test=False)
val_acc = results['accuracy']
val_loss = results['mean_outer_loss']
val_losses.append(val_loss)
val_accuracies.append(val_acc)
val_ious.append(results['iou'])
print(
f'\n validation accuracy: {val_acc}, validation loss: {val_loss}'
)
print(
'end evaluate-------------------------------------------------'
)
# Save best model
if 'accuracies_after' in results:
if (best_value is None) or (best_value <
results['accuracies_after']):
best_value = results['accuracies_after']
save_model = True
elif (best_value is None) or (best_value >
results['mean_outer_loss']):
best_value = results['mean_outer_loss']
save_model = True
else:
save_model = False
if save_model and (args.output_folder is not None):
with open(args.model_path, 'wb') as f:
torch.save(model.state_dict(), f)
print('end epoch ', epoch + 1)
elapsed_time = time.time() - start_time
print('Finished after ',
time.strftime('%H:%M:%S', time.gmtime(elapsed_time)))
r = {}
r['train_losses'] = train_losses
r['train_accuracies'] = train_accuracies
r['train_ious'] = train_ious
r['val_losses'] = val_losses
r['val_accuracies'] = val_accuracies
r['val_ious'] = val_ious
r['time'] = time.strftime('%H:%M:%S', time.gmtime(elapsed_time))
with open(os.path.join(output_folder, 'train_results.json'), 'w') as g:
json.dump(r, g)
logging.info('Saving results dict in `{0}`'.format(
os.path.abspath(os.path.join(output_folder,
'train_results.json'))))
# Plot results
plot_errors(args.num_epochs,
train_losses,
val_losses,
val_step_size=args.val_step_size,
output_folder=output_folder,
save=True,
bce_dice_focal=bce_dice_focal)
plot_accuracy(args.num_epochs,
train_accuracies,
val_accuracies,
val_step_size=args.val_step_size,
output_folder=output_folder,
save=True)
plot_iou(args.num_epochs,
train_ious,
val_ious,
val_step_size=args.val_step_size,
output_folder=output_folder,
save=True)
if hasattr(meta_train_dataset, 'close'):
meta_train_dataset.close()
meta_val_dataset.close()
