def train(name, dataset, epochs, batch_size, learning_rate, regularization,
embedding_dims, embedding_type):
dirname, _ = os.path.split(os.path.abspath(__file__))
run_uid = datetime.datetime.today().strftime('%Y-%m-%dT%H:%M:%S')
logger = StatsLogger(dirname, 'stats', name, run_uid)
print('Loading data')
X_train, y_train = load('{}_train'.format(dataset))
X_valid, y_valid = load('{}_valid'.format(dataset))
vocab = load('{}_vocab'.format(dataset)).vocab
X_train = torch.as_tensor(X_train, dtype=torch.long)
y_train = torch.as_tensor(y_train, dtype=torch.float)
X_valid = torch.as_tensor(X_valid, dtype=torch.long)
y_valid = torch.as_tensor(y_valid, dtype=torch.float)
prev_acc = 0
model = TextCNN(dataset=dataset,
input_size=X_train.size()[1],
vocab_size=len(vocab) + 1,
embedding_dims=embedding_dims,
embedding_type=embedding_type)
print(model)
print('Parameters: {}'.format(sum([p.numel() for p in \
model.parameters() if p.requires_grad])))
print('Training samples: {}'.format(len(X_train)))
if torch.cuda.is_available():
X_train = X_train.cuda()
y_train = y_train.cuda()
X_valid = X_valid.cuda()
y_valid = y_valid.cuda()
model = model.cuda()
optimizer = optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=regularization)
criterion = nn.BCEWithLogitsLoss()
print('Starting training')
for epoch in range(epochs):
epoch_loss = []
epoch_acc = []
iters = 0
total_iters = num_batches(len(X_train), batch_size)
for i, batch in enumerate(minibatch_iter(len(X_train), batch_size)):
model.train()
X_train_batch = X_train[batch]
y_train_batch = y_train[batch]
if torch.cuda.is_available():
X_train_batch = X_train_batch.cuda()
y_train_batch = y_train_batch.cuda()
optimizer.zero_grad()
output = model(X_train_batch)
train_loss = criterion(output, y_train_batch)
train_acc = accuracy(output, y_train_batch)
epoch_loss.append(train_loss.item())
epoch_acc.append(train_acc.item())
train_loss.backward()
optimizer.step()
model.eval()
train_loss, train_acc = np.mean(epoch_loss), np.mean(epoch_acc)
valid_loss, valid_acc, _ = compute_dataset_stats(
X_valid, y_valid, model, nn.BCEWithLogitsLoss(), 256)
stats = [epoch + 1, train_loss, train_acc, valid_loss, valid_acc]
epoch_string = '* Epoch {}: t_loss={:.3f}, t_acc={:.3f}, ' + \
'v_loss={:.3f}, v_acc={:.3f}'
print(epoch_string.format(*stats))
logger.write(stats)
# checkpoint model
if prev_acc < valid_acc:
prev_acc = valid_acc
model_path = os.path.join(dirname, 'checkpoints', name)
torch.save(model.state_dict(), model_path)
logger.close()
dropout_p=0.6,
eye=args.flat)
if args.cascaded_step2:
label_model_pretrained = torch.load(
args.pretrained_label_model)['label_model']
label_model.load_state_dict(label_model_pretrained)
if args.flat or args.cascaded_step2:
for param in label_model.parameters():
param.require_grad = False
doc_model = nn.DataParallel(doc_model)
label_model = nn.DataParallel(label_model)
doc_model = doc_model.cuda()
label_model = label_model.cuda()
# Loss and optimizer
criterion = Loss(use_geodesic=args.joint,
_lambda=args.geodesic_lambda,
only_label=args.cascaded_step1)
optimizer = torch.optim.Adam([{
'params': doc_model.parameters(),
'lr': doc_lr
}, {
'params': label_model.parameters(),
'lr': 0.001
}])