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