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))