def main():
global args, best_acc1, device
criterion = PrototypicalLoss().to(device)
cudnn.benchmark = True
runs_path = glob('runs/*')
max_len_exp = max([len(x) for x in runs_path]) + 2
print(f"|{'Experiment':^{max_len_exp}}|{'Loss':^17}|{'ACC':^17}|")
except_list = []
pl_mi = u"\u00B1"
for exp in glob('runs/*'):
checkpoint, args = None, None
files = glob(exp + '/*')
for file in files:
if file.endswith('model_best.pth'):
checkpoint = torch.load(os.path.abspath(file))
elif file.endswith('.json'):
params = json.load(open(os.path.abspath(file)))
args = SimpleNamespace(**params)
if checkpoint is None or args is None:
except_list.append(f"checkpoint and params are not exist in {exp}")
continue
if args.dataset == 'omniglot':
test_loader = get_dataloader(args, 'test')
else:
test_loader = get_dataloader(args, 'val')
input_dim = 1 if args.dataset == 'omniglot' else 3
if args.model == 'protonet':
model = ProtoNet(input_dim).to(device)
else:
model = ResNet(input_dim).to(device)
model.load_state_dict(checkpoint['model_state_dict'])
best_acc1 = checkpoint['best_acc1']
loss_list, acc_list = test(test_loader, model, criterion)
loss, loss_moe = margin_of_error(loss_list)
acc, acc_moe = margin_of_error(acc_list)
loss_string = f'{loss:.3f} {pl_mi} {loss_moe:.3f}'
acc_string = f'{acc:.3f} {pl_mi} {acc_moe:.3f}'
print(f"|{exp:^{max_len_exp}}|{loss_string:^16}|{acc_string:^16}|")
if len(except_list):
pp(except_list)
def __init__(self, args, mode='train'):
# init model and optimizer
self.model = ProtoNet(args)
print(self.model)
if torch.cuda.is_available():
self.model.cuda()
if mode == 'train':
if args.use_attention:
self.optimizer = torch.optim.Adam(
[{
'params': self.model.encoder.parameters(),
'lr': 0.0001
}, {
'params': self.model.base_learner.parameters()
}, {
'params': self.model.att_learner.parameters()
}],
lr=args.lr)
else:
self.optimizer = torch.optim.Adam(
[{
'params': self.model.encoder.parameters(),
'lr': 0.0001
}, {
'params': self.model.base_learner.parameters()
}, {
'params': self.model.linear_mapper.parameters()
}],
lr=args.lr)
#set learning rate scheduler
self.lr_scheduler = optim.lr_scheduler.StepLR(
self.optimizer, step_size=args.step_size, gamma=args.gamma)
# load pretrained model for point cloud encoding
self.model = load_pretrain_checkpoint(
self.model, args.pretrain_checkpoint_path)
elif mode == 'test':
# Load model checkpoint
self.model = load_model_checkpoint(self.model,
args.model_checkpoint_path,
mode='test')
else:
raise ValueError('Wrong GMMLearner mode (%s)! Option:train/test' %
mode)
def get_parser(parser=None):
if parser is None:
parser = argparse.ArgumentParser(description='RelationNet')
parser = ProtoNet.get_parser(parser)
parser.add_argument('--loss_type',
type=str,
choices=['mse', 'softmax'],
default='mse')
return parser
train_loader = DataLoader(dataset=trainset,
batch_sampler=train_sampler,
num_workers=8,
pin_memory=True)
valset = Dataset('val', args)
val_sampler = CategoriesSampler(valset.label, 500, args.validation_way,
args.shot + args.query)
val_loader = DataLoader(dataset=valset,
batch_sampler=val_sampler,
num_workers=8,
pin_memory=True)
model = None
if args.model_type.lower() == 'protonet':
model = ProtoNet(args)
elif args.model_type.lower() == 'hypnet':
model = HypNet(args)
elif args.model_type.lower() == 'protonetwithhyperbolic':
model = ProtoNetWithHyperbolic(args)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
if args.lr_decay:
lr_scheduler = torch.optim.lr_scheduler.StepLR(
optimizer, step_size=args.step_size, gamma=args.gamma)
# load pre-trained model (no FC weights)
model_dict = model.state_dict()
if args.init_weights is not None:
pretrained_dict = torch.load(args.init_weights)['params']
hidden_size=args.lang_hidden_size,
rnn=args.rnn_type,
num_layers=args.rnn_num_layers,
dropout=args.rnn_dropout,
)
l3_model = l3_model.cuda()
embedding_model = embedding_model.cuda()
lang_model = lang_model.cuda()
model = ProtoNet(
model_dict[args.model],
n_way=args.test_n_way,
n_support=args.n_shot,
# Language options
lsl=args.lsl,
language_model=lang_model,
lang_supervision=args.lang_supervision,
l3=args.l3,
l3_model=l3_model,
l3_n_infer=args.l3_n_infer,
)
model = model.cuda()
if args.save_iter != -1:
modelfile = get_assigned_file(args.checkpoint_dir, args.save_iter)
else:
modelfile = get_best_file(args.checkpoint_dir)
if modelfile is not None:
tmp = torch.load(modelfile)
def main():
global args, best_acc1, device
# Init seed
np.random.seed(args.manual_seed)
torch.manual_seed(args.manual_seed)
torch.cuda.manual_seed(args.manual_seed)
if args.dataset == 'omniglot':
train_loader, val_loader = get_dataloader(args, 'trainval', 'test')
input_dim = 1
else:
train_loader, val_loader = get_dataloader(args, 'train', 'val')
input_dim = 3
if args.model == 'protonet':
model = ProtoNet(input_dim).to(device)
print("ProtoNet loaded")
else:
model = ResNet(input_dim).to(device)
print("ResNet loaded")
criterion = PrototypicalLoss().to(device)
optimizer = torch.optim.Adam(model.parameters(), args.lr)
cudnn.benchmark = True
if args.resume:
try:
checkpoint = torch.load(
sorted(glob(f'{args.log_dir}/checkpoint_*.pth'), key=len)[-1])
except Exception:
checkpoint = torch.load(args.log_dir + '/model_best.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
print(f"load checkpoint {args.exp_name}")
else:
start_epoch = 1
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer=optimizer,
gamma=args.lr_scheduler_gamma,
step_size=args.lr_scheduler_step)
print(
f"model parameter : {sum(p.numel() for p in model.parameters() if p.requires_grad)}"
)
for epoch in range(start_epoch, args.epochs + 1):
train_loss = train(train_loader, model, optimizer, criterion, epoch)
is_test = False if epoch % args.test_iter else True
if is_test or epoch == args.epochs or epoch == 1:
val_loss, acc1 = validate(val_loader, model, criterion, epoch)
if acc1 >= best_acc1:
is_best = True
best_acc1 = acc1
else:
is_best = False
save_checkpoint(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer_state_dict': optimizer.state_dict(),
}, is_best, args)
if is_best:
writer.add_scalar("BestAcc", acc1, epoch)
print(
f"[{epoch}/{args.epochs}] {train_loss:.3f}, {val_loss:.3f}, {acc1:.3f}, # {best_acc1:.3f}"
)
else:
print(f"[{epoch}/{args.epochs}] {train_loss:.3f}")
scheduler.step()
writer.close()
)
print("Loading val data\n")
val_loader = val_datamgr.get_data_loader(
val_file,
aug=False,
lang_dir=constants.LANG_DIR,
normalize=True,
vocab=vocab,
)
# a batch for SetDataManager: a [n_way, n_support + n_query, dim, w, h] tensor
model = ProtoNet(
model_dict[args.model],
**train_few_shot_args,
# Language options
lsl=args.lsl,
language_model=lang_model,
lang_supervision=args.lang_supervision,
l3=args.l3,
l3_model=l3_model,
l3_n_infer=args.l3_n_infer)
model = model.cuda()
os.makedirs(args.checkpoint_dir, exist_ok=True)
start_epoch = args.start_epoch
stop_epoch = args.stop_epoch
if args.resume:
resume_file = get_resume_file(args.checkpoint_dir)
if resume_file is not None:
class ProtoLearner(object):
def __init__(self, args, mode='train'):
# init model and optimizer
self.model = ProtoNet(args)
print(self.model)
if torch.cuda.is_available():
self.model.cuda()
if mode == 'train':
if args.use_attention:
self.optimizer = torch.optim.Adam(
[{
'params': self.model.encoder.parameters(),
'lr': 0.0001
}, {
'params': self.model.base_learner.parameters()
}, {
'params': self.model.att_learner.parameters()
}],
lr=args.lr)
else:
self.optimizer = torch.optim.Adam(
[{
'params': self.model.encoder.parameters(),
'lr': 0.0001
}, {
'params': self.model.base_learner.parameters()
}, {
'params': self.model.linear_mapper.parameters()
}],
lr=args.lr)
#set learning rate scheduler
self.lr_scheduler = optim.lr_scheduler.StepLR(
self.optimizer, step_size=args.step_size, gamma=args.gamma)
# load pretrained model for point cloud encoding
self.model = load_pretrain_checkpoint(
self.model, args.pretrain_checkpoint_path)
elif mode == 'test':
# Load model checkpoint
self.model = load_model_checkpoint(self.model,
args.model_checkpoint_path,
mode='test')
else:
raise ValueError('Wrong GMMLearner mode (%s)! Option:train/test' %
mode)
def train(self, data):
"""
Args:
data: a list of torch tensors wit the following entries.
- support_x: support point clouds with shape (n_way, k_shot, in_channels, num_points)
- support_y: support masks (foreground) with shape (n_way, k_shot, num_points)
- query_x: query point clouds with shape (n_queries, in_channels, num_points)
- query_y: query labels with shape (n_queries, num_points)
"""
[support_x, support_y, query_x, query_y] = data
self.model.train()
query_logits, loss = self.model(support_x, support_y, query_x, query_y)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
self.lr_scheduler.step()
query_pred = F.softmax(query_logits, dim=1).argmax(dim=1)
correct = torch.eq(query_pred,
query_y).sum().item() # including background class
accuracy = correct / (query_y.shape[0] * query_y.shape[1])
return loss, accuracy
def test(self, data):
"""
Args:
support_x: support point clouds with shape (n_way, k_shot, in_channels, num_points)
support_y: support masks (foreground) with shape (n_way, k_shot, num_points), each point \in {0,1}.
query_x: query point clouds with shape (n_queries, in_channels, num_points)
query_y: query labels with shape (n_queries, num_points), each point \in {0,..., n_way}
"""
[support_x, support_y, query_x, query_y] = data
self.model.eval()
with torch.no_grad():
logits, loss = self.model(support_x, support_y, query_x, query_y)
pred = F.softmax(logits, dim=1).argmax(dim=1)
correct = torch.eq(pred, query_y).sum().item()
accuracy = correct / (query_y.shape[0] * query_y.shape[1])
return pred, loss, accuracy