def main(args):
seed = 2020
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# Make logdir
if not os.path.exists(args.checkpoints_dir):
os.makedirs(args.checkpoints_dir)
# Load dataset
train_dataloader = get_dataloader('train', args.bs, True, args.nw)
val_dataloader = get_dataloader('val', args.bs, False, args.nw)
# Model
model = SimpleModel()
optimizer = torch.optim.SGD(
model.parameters(), lr=args.lr, momentum=0.9,
weight_decay=args.wd)
model.cuda()
train(args, train_dataloader, val_dataloader,
model, optimizer)
def main(args):
# Load dataset
test_dataloader = get_dataloader('test', args.bs, False, args.nw)
# Model
model = SimpleModel()
model.cuda()
ckpt = torch.load(os.path.join(args.checkpoints_dir, 'last_ckpt.pth'))
model.load_state_dict(ckpt['model_state'])
result = test(args, test_dataloader, model)
# Make csv file
df = pd.DataFrame({'id': test_dataloader.dataset.ids, 'category': result})
df.to_csv('out.csv', index=False)