def main(parser, logger):
print('--> Preparing Dataset:')
trainset = OmniglotDataset(mode='train', root=parser.dataset_root)
trainloader = data.DataLoader(trainset, batch_size=100, shuffle=True, num_workers=0)
valset = dataloader(parser, 'val')
testset = dataloader(parser, 'test')
print('--> Building Model:')
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
model = Network.resnet18().to(device)
model = DataParallel(model)
metric = ArcMarginProduct(256, len(np.unique(trainset.y)), s=30, m=0.5).to(device)
metric = DataParallel(metric)
criterion = torch.nn.CrossEntropyLoss()
print('--> Initializing Optimizer and Scheduler:')
optimizer = torch.optim.Adam(
[{'params':model.parameters(), 'weight_decay':5e-4},
{'params':[metric.weight], 'weight_decay':5e-4}],
lr=parser.learning_rate, weight_decay=0.0005)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optimizer,
gamma=parser.lr_scheduler_gamma,
step_size=parser.lr_scheduler_step)
best_acc = 0
best_state = model.state_dict()
for epoch in range(parser.epochs):
print('\nEpoch: %d' % epoch)
# Training
train_loss = 0
train_acc = 0
train_correct = 0
train_total = 0
model.train()
for batch_index, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device).long()
feature = model(inputs)
output = metric(feature, targets)
loss = criterion(output, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss.item()
_, predicted = output.max(1)
train_total += targets.size(0)
train_correct += predicted.eq(targets).sum().item()
scheduler.step()
train_acc = 100.*train_correct / train_total
print('Training Loss: {} | Accuracy: {}'.format(train_loss/train_total, train_acc))
# Validating
val_correct = 0
val_total = 0
model.eval()
for batch_index, (inputs, targets) in enumerate(valset):
inputs = inputs.to(device)
targets = targets.to(device)
feature = model(inputs)
correct = eval(input=feature, target=targets, n_support=parser.num_support_val)
val_correct += correct
val_total += parser.classes_per_it_val * parser.num_query_val
val_acc = 100.*val_correct / val_total
print('Validating Accuracy: {}'.format(val_acc))
if val_acc > best_acc:
best_acc = val_acc
best_state = model.state_dict()
test_correct = 0
test_total = 0
model.load_state_dict(best_state)
for epoch in range(10):
for batch_index, (inputs, targets) in enumerate(testset):
inputs = inputs.to(device)
targets = targets.to(device)
feature = model(inputs)
correct = eval(input=feature, target=targets, n_support=parser.num_support_val)
test_correct += correct
test_total += parser.classes_per_it_val * parser.num_query_val
test_acc = 100. * test_correct / test_total
print('Testing Accuracy: {}'.format(test_acc))
def main(parser, logger):
print('--> Preparing Dataset:')
trainset = dataloader(parser, 'train')
valset = dataloader(parser, 'val')
testset = dataloader(parser, 'test')
print('--> Preparing Word Embedding Model')
print('--> Building Model:')
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
model = resnet18().to(device)
print('--> Initializing Optimizer and Scheduler')
optimizer = torch.optim.Adam(params=model.parameters(),
lr=parser.learning_rate)
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer=optimizer,
gamma=parser.lr_scheduler_gamma,
step_size=parser.lr_scheduler_step)
train_loss = []
train_acc = []
val_loss = []
val_acc = []
best_acc = 0
best_model_path = os.path.join(parser.experiment_root, 'best_model.pth')
best_state = model.state_dict()
for epoch in range(parser.epochs):
print('\nEpoch: %d' % epoch)
# Training
model.train()
for batch_index, (inputs, targets,
targets_embedding) in enumerate(trainset):
support_samples, support_embeddings, query_samples, query_embeddings = \
splitInput(inputs, targets, targets_embedding, parser.num_support_tr)
support_samples = support_samples.to(device)
query_samples = query_samples.to(device)
targets = targets.to(device)
support_embeddings = support_embeddings.to(device)
query_embeddings = query_embeddings.to(device)
support_output = model(support_samples,
support_embeddings,
mode='support')
query_output = model(query_samples, query_embeddings, mode='query')
loss, acc = loss_fn2(support_input=support_output,
query_input=query_output,
targets=targets,
n_support=parser.num_support_tr)
train_loss.append(loss.item())
train_acc.append(acc.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
avg_loss = np.mean(train_loss[-parser.iterations:])
avg_acc = 100. * np.mean(train_acc[-parser.iterations:])
print('Training Loss: {} | Accuracy: {}'.format(avg_loss, avg_acc))
# Validating
model.eval()
for batch_index, (inputs, targets,
targets_embedding) in enumerate(valset):
support_samples, support_embeddings, query_samples, query_embeddings = \
splitInput(inputs, targets, targets_embedding, parser.num_support_tr)
support_samples = support_samples.to(device)
query_samples = query_samples.to(device)
targets = targets.to(device)
support_embeddings = support_embeddings.to(device)
query_embeddings = query_embeddings.to(device)
support_output = model(support_samples,
support_embeddings,
mode='support')
query_output = model(query_samples, query_embeddings, mode='query')
loss, acc = loss_fn2(support_input=support_output,
query_input=query_output,
targets=targets,
n_support=parser.num_support_val)
val_loss.append(loss.item())
val_acc.append(acc.item())
avg_loss = np.mean(val_loss[-parser.iterations:])
avg_acc = 100. * np.mean(val_acc[-parser.iterations:])
print('Validating Loss: {} | Accuracy: {}'.format(avg_loss, avg_acc))
info = {'val_loss': avg_loss, 'val_accuracy': avg_acc}
for tag, value in info.items():
logger.scalar_summary(tag, value, epoch + 1)
if avg_acc > best_acc:
torch.save(model.state_dict(), best_model_path)
best_acc = avg_acc
best_state = model.state_dict()
# Testing
model.load_state_dict(best_state)
test_acc = []
for epoch in range(10):
for batch_index, (inputs, targets,
targets_embedding) in enumerate(testset):
support_samples, support_embeddings, query_samples, query_embeddings = \
splitInput(inputs, targets, targets_embedding, parser.num_support_tr)
support_samples = support_samples.to(device)
query_samples = query_samples.to(device)
targets = targets.to(device)
support_embeddings = support_embeddings.to(device)
query_embeddings = query_embeddings.to(device)
support_output = model(support_samples,
support_embeddings,
mode='support')
query_output = model(query_samples, query_embeddings, mode='query')
_, acc = loss_fn2(support_input=support_output,
query_input=query_output,
targets=targets,
n_support=parser.num_support_tr)
test_acc.append(acc.item())
avg_acc = 100. * np.mean(test_acc)
logger.scalar_summary('test_accuracy', avg_acc, 1)
print('*****Testing Accuracy: {}'.format(avg_acc))