criterion = nn.CrossEntropyLoss()
# Scheduler
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1.5 * 1e-3, steps_per_epoch=len(train_dataloader),
epochs=num_epochs)
# Training
# optim_gen = lambda parameters, lr: SGD(parameters, lr=lr)
# find_lr(model, optim_gen, min_lr, max_lr, num_epochs, train_dataloader, val_dataloader, criterion, device, batch_size,
# batches_per_epoch, comet_experiment)
save_path = pathlib.Path('models') / name
save_path.mkdir(parents=True, exist_ok=True)
trainer = Trainer(model, train_dataloader, val_dataloader, criterion, optimizer, None, device, TRAFFIC_LABELS,
num_epochs, batch_size, batches_per_epoch, comet_experiment, save_path)
try:
trainer.fit()
except KeyboardInterrupt:
pass
# Prediction
data_path_root_test = pathlib.Path('test/')
test_anno = pd.DataFrame({'id': [f'pic{num:06}' for num in range(10699)]})
test_dataset = MyDataset(data_dir=data_path_root_test, data_anno=test_anno, phase='test')
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
preds = trainer.predict(test_dataloader)
submit = pd.DataFrame({'id': [f'pic{num:06}' for num in range(10699)],
'category': [TRAFFIC_LABELS[pred] for pred in preds]})
submit.to_csv(save_path / 'submit.csv')
print('Done')
if not args.load_pretrained_path:
model = TransformerMaskPredict(config)
else:
model = TransformerMaskPredict.from_pretrained(
args.load_pretrained_path)
tokenizer = Tokenizer(model.config.encoder_id, model.config.decoder_id,
model.config.length_token)
dl = DataLoader(args, tokenizer)
tr_dataset, val_dataset, test_dataset = dl()
trainer = Trainer(model, tokenizer, args)
if args.load_training_state:
trainer.load_training_state_dict(args.base_dir)
trainer.fit(tr_dataset, val_dataset)
if args.save_pretrained_path:
trainer.model.save_pretrained(
os.path.join(args.base_dir, args.save_pretrained_path))
tst_loss = trainer.evaluate(test_dataset)
wandb.log({"tst_loss": tst_loss})
for mode, dataset in zip(["tr", "val", "tst"],
[tr_dataset, val_dataset, test_dataset]):
out = fetch_translations_and_bleu(model, dataset, tokenizer,
args.iterations, args.B,
args.bleu_num_samples)
data = list(zip(out["src"], out["tgt"], out["pred"]))
