def load_model(lr):
model = ConvolutionalNeuralNetwork(feat_train.shape[1])
loss_fnc = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr = lr, momentum = 0.9)
#optimizer = torch.optim.Adam(model.parameters(), lr = lr)
return model, loss_fnc, optimizer
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`.')
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 = ConvolutionalNeuralNetwork(1,
args.num_ways,
hidden_size=args.hidden_size)
model.to(device=args.device)
model.train()
meta_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)
test_inputs, test_targets = batch['test']
test_inputs = test_inputs.to(device=args.device)
test_targets = test_targets.to(device=args.device)
outer_loss = torch.tensor(0., device=args.device)
accuracy = torch.tensor(0., device=args.device)
for task_idx, (train_input, train_target, test_input,
test_target) in enumerate(
zip(train_inputs, train_targets, test_inputs,
test_targets)):
train_logit = model(train_input)
inner_loss = F.cross_entropy(train_logit, train_target)
model.zero_grad()
params = gradient_update_parameters(
model,
inner_loss,
step_size=args.step_size,
first_order=args.first_order)
test_logit = model(test_input, params=params)
outer_loss += F.cross_entropy(test_logit, test_target)
with torch.no_grad():
accuracy += get_accuracy(test_logit, test_target)
outer_loss.div_(args.batch_size)
accuracy.div_(args.batch_size)
outer_loss.backward()
meta_optimizer.step()
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, 'maml_omniglot_'
'{0}shot_{1}way.th'.format(args.num_shots, args.num_ways))
with open(filename, 'wb') as f:
state_dict = model.state_dict()
torch.save(state_dict, f)
def train(args):
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 = ConvolutionalNeuralNetwork(1,
args.num_ways,
hidden_size=args.hidden_size)
model.to(device=args.device)
model.train()
meta_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)
test_inputs, test_targets = batch['test']
test_inputs = test_inputs.to(device=args.device)
test_targets = test_targets.to(device=args.device)
outer_loss = torch.tensor(0., device=args.device)
accuracy = torch.tensor(0., device=args.device)
for task_idx, (train_input, train_target, test_input,
test_target) in enumerate(
zip(train_inputs, train_targets, test_inputs,
test_targets)):
train_logit = model(train_input)
inner_loss = F.cross_entropy(train_logit, train_target)
model.zero_grad()
params = update_parameters(model,
inner_loss,
step_size=args.step_size,
first_order=args.first_order)
test_logit = model(test_input, params=params)
outer_loss += F.cross_entropy(test_logit, test_target)
with torch.no_grad():
accuracy += get_accuracy(test_logit, test_target)
outer_loss.div_(args.batch_size)
accuracy.div_(args.batch_size)
outer_loss.backward()
meta_optimizer.step()
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, 'maml_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)
args.out_channels, _ = get_outputs_c_h(args.backbone, args.image_len)
model = ConvolutionalNeuralNetwork(args.backbone, args.out_channels, args.num_ways)
teacher_model = ConvolutionalNeuralNetwork(args.teacher_backbone, args.out_channels, args.num_ways)
if args.use_cuda:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
num_gpus = torch.cuda.device_count()
if args.multi_gpu:
model = DataParallel(model)
teacher_model = DataParallel(teacher_model)
model = model.cuda()
teacher_model = teacher_model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=0.0005)
# download teacher trained model
if args.teacher_resume and args.teacher_resume_folder is not None:
teacher_model_path = os.path.join(args.teacher_resume_folder, ('_'.join(['teacher', args.train_data, args.test_data, args.teacher_backbone, 'max_acc']) + '.pt'))
teacher_state = torch.load(teacher_model_path)
if args.multi_gpu:
teacher_model.module.load_state_dict(teacher_state)
else:
teacher_model.load_state_dict(teacher_state)
# training from the checkpoint
if args.resume and args.resume_folder is not None:
# load checkpoint
checkpoint_path = os.path.join(args.resume_folder, ('_'.join([args.model_name, args.train_data, args.test_data, args.backbone, 'max_acc']) + '_checkpoint.pt.tar')) # tag='max_acc' can be changed
state = torch.load(checkpoint_path)
def train(args):
dataset = miniimagenet(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 = ConvolutionalNeuralNetwork(3,
84,
args.num_ways,
hidden_size=args.hidden_size)
model.to(device=args.device)
model.train()
meta_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)
test_inputs, test_targets = batch['test']
test_inputs = test_inputs.to(device=args.device)
test_targets = test_targets.to(device=args.device)
outer_loss = torch.tensor(0., device=args.device)
accuracy = torch.tensor(0., device=args.device)
for task_idx, (train_input, train_target, test_input,
test_target) in enumerate(
zip(train_inputs, train_targets, test_inputs,
test_targets)):
train_logit = model(train_input)
inner_loss = F.cross_entropy(train_logit, train_target)
# writer.add_scalar('Loss/inner_loss', np.random.random(), task_idx)
grid = torchvision.utils.make_grid(train_input)
writer.add_image('images', grid, 0)
writer.add_graph(model, train_input)
model.zero_grad()
params = update_parameters(model,
inner_loss,
step_size=args.step_size,
first_order=args.first_order)
test_logit = model(test_input, params=params)
outer_loss += F.cross_entropy(test_logit, test_target)
# writer.add_scalar('Loss/outer_loss', np.random.random(), n_iter)
for name, grads in model.meta_named_parameters():
writer.add_histogram(name, grads, batch_idx)
with torch.no_grad():
accuracy += get_accuracy(test_logit, test_target)
writer.add_histogram('meta parameters', grads, batch_idx)
outer_loss.div_(args.batch_size)
accuracy.div_(args.batch_size)
outer_loss.backward()
meta_optimizer.step()
pbar.set_postfix(accuracy='{0:.4f}'.format(accuracy.item()))
writer.add_scalar('Accuracy/test', accuracy.item(), batch_idx)
if batch_idx >= args.num_batches:
break
writer.close()
# Save model
if args.output_folder is not None:
filename = os.path.join(
args.output_folder, 'maml_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)
# construct model and optimizer
args.image_len = 28 if args.train_data == 'omniglot' else 84
args.out_channels, _ = get_outputs_c_h(args.backbone, args.image_len)
model = ConvolutionalNeuralNetwork(args.backbone, args.out_channels,
args.num_ways)
if args.use_cuda:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
num_gpus = torch.cuda.device_count()
if args.multi_gpu:
model = DataParallel(model)
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=0.0005)
# training from the checkpoint
if args.resume and args.resume_folder is not None:
# load checkpoint
checkpoint_path = os.path.join(args.resume_folder, ('_'.join([
args.model_name, args.train_data, args.test_data, args.backbone,
'max_acc'
]) + '_checkpoint.pt.tar')) # tag='max_acc' can be changed
state = torch.load(checkpoint_path)
if args.multi_gpu:
model.module.load_state_dict(state['model'])
else:
model.load_state_dict(state['model'])
def load_model(lr, input_size, output_size):
model = ConvolutionalNeuralNetwork(input_size, output_size)
loss_fnc = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
return model, loss_fnc, optimizer