Spaces:
Runtime error
Runtime error
File size: 4,257 Bytes
7900c16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
"""
This script provides an example to wrap TencentPretrain for classification with grid search.
"""
import sys
import os
import torch
import random
import argparse
from itertools import product
tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(tencentpretrain_dir)
from tencentpretrain.opts import *
from finetune.run_classifier import *
def main():
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
finetune_opts(parser)
tokenizer_opts(parser)
parser.add_argument("--soft_targets", action='store_true',
help="Train model with logits.")
parser.add_argument("--soft_alpha", type=float, default=0.5,
help="Weight of the soft targets loss.")
parser.add_argument("--batch_size_list", default=[32, 64], nargs='+', type=int,
help="A list of batch sizes for grid search.")
parser.add_argument("--learning_rate_list", type=float, default=[3e-5, 1e-4, 3e-4], nargs='+',
help="A list of learning rate for grid search.")
parser.add_argument("--epochs_num_list", type=int, default=[3, 5, 8], nargs='+',
help="A list of number of epochs for grid search.")
adv_opts(parser)
args = parser.parse_args()
# Load the hyperparameters from the config file.
args = load_hyperparam(args)
set_seed(args.seed)
# Count the number of labels.
args.labels_num = count_labels_num(args.train_path)
# Build tokenizer.
args.tokenizer = str2tokenizer[args.tokenizer](args)
# Get logger.
args.logger = init_logger(args)
best_acc = 0
config = {}
# Build dataset
trainset = read_dataset(args, args.train_path)
instances_num = len(trainset)
for batch_size, learning_rate, epochs_num in product(args.batch_size_list, args.learning_rate_list, args.epochs_num_list):
args.learning_rate = learning_rate
args.batch_size = batch_size
args.epochs_num = epochs_num
args.train_steps = int(instances_num * args.epochs_num / batch_size) + 1
# Build classification model.
model = Classifier(args)
args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(args.device)
load_or_initialize_parameters(args, model)
optimizer, scheduler = build_optimizer(args, model)
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)
args.amp = amp
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
args.model = model
if args.use_adv:
args.adv_method = str2adv[args.adv_type](model)
# Training phase.
total_loss, _, _ = 0., 0., 0.
for _ in range(1, args.epochs_num + 1):
random.shuffle(trainset)
src = torch.LongTensor([example[0] for example in trainset])
tgt = torch.LongTensor([example[1] for example in trainset])
seg = torch.LongTensor([example[2] for example in trainset])
if args.soft_targets:
soft_tgt = torch.FloatTensor([example[3] for example in trainset])
else:
soft_tgt = None
model.train()
for i, (src_batch, tgt_batch, seg_batch, soft_tgt_batch) in enumerate(batch_loader(batch_size, src, tgt, seg, soft_tgt)):
_ = train_model(args, model, optimizer, scheduler, src_batch, tgt_batch, seg_batch, soft_tgt_batch)
acc, _ = evaluate(args, read_dataset(args, args.dev_path))
if acc > best_acc:
best_acc = acc
config = {"learning_rate": learning_rate, "batch_size": batch_size, "epochs_num": epochs_num}
args.logger.info("On configuration: {}.\n".format(config))
args.logger.info("Best Acc. is: {:.4f}, on configuration {}.".format(best_acc, config))
if __name__ == "__main__":
main()
|