sampler=train_set_sampler,
num_workers=8)
val_loader = DataLoader(dataset,
batch_size=batch_size,
sampler=validating_set_sampler,
num_workers=8)
# TensorboardX writer
writer = SummaryWriter(main_directory + '/runs/' + experiment_name +
'_' + str(seed))
# The model
torch.manual_seed(0)
model = WideNet(len(class_names)).to(device)
logger.info('Net parameters number : {}'.format(
utils.compute_total_parameter_number(model)))
optimizer = optim.SGD(model.parameters(),
lr=0.05,
momentum=0.9,
weight_decay=0.001)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[50, 100, 150],
gamma=0.1)
# Train and test
logger.info('Start training')
for epoch in range(1, max_epochs + 1):
gpu_map = utils.get_gpu_usage_map(0)
logger.info('GPU usage : {}'.format(gpu_map))
train(model, device, train_loader, optimizer, epoch, writer=writer)
indices = np.where(labels == cls)
indices = np.random.permutation(indices[0])
train_indices.append(indices[:int(len(indices) * (1 - validating_ratio))])
val_indices.append(indices[int(len(indices) * (1 - validating_ratio)):])
train_set_sampler = sampler.SubsetRandomSampler(np.concatenate(train_indices))
validating_set_sampler = sampler.SubsetRandomSampler(np.concatenate(val_indices))
train_loader = DataLoader(dataset, batch_size=batch_size, sampler=train_set_sampler, num_workers=8)
val_loader = DataLoader(dataset, batch_size=batch_size, sampler=validating_set_sampler, num_workers=8)
# TensorboardX writer
writer = SummaryWriter(main_directory + '/runs/' + experiment_name + '_' + str(seed))
# The model
torch.manual_seed(0)
model = WideNetMasked(len(class_names)).to(device)
logger.info('Net parameters number : {}'.format(utils.compute_total_parameter_number(model)))
optimizer = optim.SGD(model.parameters(), lr=0.05, momentum=0.9, weight_decay=0.001)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[50, 100, 150], gamma=0.1)
# Train and test
logger.info('Start training')
for epoch in range(1, max_epochs + 1):
gpu_map = utils.get_gpu_usage_map(0)
logger.info('GPU usage : {}'.format(gpu_map))
train(model, device, train_loader, optimizer, epoch, writer=writer)
test(model, device, val_loader, epoch, writer=writer)
scheduler.step(epoch=epoch)
if epoch % 100 == 0:
torch.save(model.state_dict(), experiment_directory + '/model_' + str(seed) + '_epoch_' + str(epoch))
