num_classes=2)
train_dataset = ConcatDataset(
[train_dataset_augmented, train_dataset_base])
val_dataset = create_dataset(val['X'], val['y'], model_name, 256, mix=None)
train_dataloader = DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=3,
shuffle=True)
val_dataloader = DataLoader(val_dataset,
batch_size=64,
num_workers=3,
shuffle=False)
model = MixText(2, True).cuda()
model = nn.DataParallel(model)
wandb.watch(model, log=all)
optimizer = AdamW([
{
"params": model.module.bert.parameters(),
"lr": args.lrmain
},
{
"params": model.module.linear.parameters(),
"lr": args.lrlast
},
])
decayRate = args.lr_decay