def get_optimizer(net, state=None):
optimizer = torch.optim.Adam(net.parameters(),
lr=args.lr,
betas=(0, 0.999))
if state is not None:
optimizer.load_state_dict(state)
return optimizer
def set_learning_rate(optimizer, lr):
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Build model, optimizer, and set states
meta_net = OmniglotModel(args.classes)
if args.cuda:
meta_net.cuda()
meta_optimizer = torch.optim.SGD(meta_net.parameters(), lr=args.meta_lr)
info = {}
state = None
# checkpoint is directory -> Find last model or '' if does not exist
if os.path.isdir(args.checkpoint):
latest_checkpoint = find_latest_file(check_dir)
if latest_checkpoint:
print 'Latest checkpoint found:', latest_checkpoint
args.checkpoint = os.path.join(check_dir, latest_checkpoint)
else:
args.checkpoint = ''
return train_loss, test_loss, acc
fc_dim = 32
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
data_dir = os.path.join(parent_dir, 'data')
word2index_path = os.path.join(data_dir, 'word2index.pkl')
emb_mat_path = os.path.join(data_dir, 'emb_mat.npy')
word2index = load_pkl(word2index_path)
emb_mat = np.load(emb_mat_path)
vocab_size = len(word2index) + 2
# Build model, optimizer, and set states
meta_net = OmniglotModel(fc_dim, args.classes, vocab_size, emb_mat, args.cuda)
meta_net.cuda(args.cuda)
meta_optimizer = torch.optim.SGD(meta_net.parameters(), lr=args.meta_lr)
state = None
cca_state = None
# fast_optimizer
cca_optimizer = torch.optim.RMSprop(meta_net.parameters(),
lr=1e-3,
weight_decay=1e-5)
outdim_size = fc_dim # 32
use_all_singular_values = False
loss_fn = nn.CrossEntropyLoss()
cca_loss = CCALoss(outdim_size, use_all_singular_values, args.cuda)