def _save_model_components(path: str, input_processor: InputProcessor,
bertram_config: bertram.BertramConfig,
model: bertram.Bertram, meta_info: Dict) -> None:
output_ip_file = os.path.join(path, bertram.IP_NAME)
input_processor.save(output_ip_file)
output_bc_file = os.path.join(path, bertram.CONFIG_NAME)
bertram_config.save(output_bc_file)
model_to_save = model.module if hasattr(model, 'module') else model
output_model_file = os.path.join(path, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(path, CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
meta_file = os.path.join(path, META_NAME)
with open(meta_file, 'w') as f:
for k, v in meta_info.items():
f.write('{}: {}\n'.format(k, v))
def main(args):
parser = argparse.ArgumentParser("Train a new BERTRAM instance")
# Required parameters
parser.add_argument('--model_cls', default='bert', choices=['bert', 'roberta'],
help="The transformer model class, either 'bert' or 'roberta'.")
parser.add_argument("--bert_model", default=None, type=str, required=True,
help="The pretrained model to use (e.g., 'bert-base-uncased'.")
parser.add_argument("--output_dir", default=None, type=str, required=True,
help="The output directory where the model checkpoints will be written.")
parser.add_argument('--train_dir', type=str, required=True,
help="The directory in which the buckets for training are stored. "
"Each bucket should be a text file containing lines of the form "
"<WORD><TAB><CONTEXT_1><TAB>...<CONTEXT_n>")
parser.add_argument('--vocab', type=str, required=True,
help="The file in which the vocabulary to be used for training is stored."
"Each line should be of the form <WORD> <COUNT>")
parser.add_argument('--emb_file', type=str, required=True,
help="The file in which the target embeddings for mimicking are stored.")
parser.add_argument('--emb_dim', type=int, required=True,
help="The number of dimensions for the target embeddings.")
parser.add_argument('--mode', choices=bertram.MODES, required=True,
help="The BERTRAM mode (e.g., 'form', 'add', 'replace').")
# Other parameters
parser.add_argument("--cache_dir", default="", type=str,
help="Where to store the pre-trained models downloaded from s3")
parser.add_argument("--max_seq_length", default=128, type=int,
help="The maximum total input sequence length after WordPiece tokenization.")
parser.add_argument("--train_batch_size", default=32, type=int,
help="Total batch size for training.")
parser.add_argument("--learning_rate", default=5e-5, type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs", default=3.0, type=float,
help="Total number of training epochs to perform.")
parser.add_argument("--warmup_proportion", default=0.1, type=float,
help="Proportion of training to perform linear learning rate warmup for.")
parser.add_argument('--seed', type=int, default=42,
help="Random seed for initialization")
parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
help="Number of updates steps to accumulate before performing a backward/update pass.")
parser.add_argument('--save_epochs', type=int, nargs='+', default=[1, 5, 10],
help="The number of epochs after which a checkpoint is saved.")
parser.add_argument('--min_word_count', '-mwc', type=int, default=100,
help="The minimum number of occurrences for a word to be used as training target.")
parser.add_argument('--num_buckets', type=int, default=25,
help="The number of buckets in the training directory.")
parser.add_argument('--emb_format', type=str, choices=['text', 'gensim'], default='text',
help="The format in which target embeddings are stored.")
parser.add_argument('--no_finetuning', action='store_true',
help="Whether not to finetune the underlying transformer language model.")
parser.add_argument('--optimize_only_combinator', action='store_true',
help="Whether to freeze both the underyling transformer language model and the ngram"
"embeddings during training.")
# Context parameters
parser.add_argument('--smin', type=int, default=20,
help="The minimum number of contexts per word.")
parser.add_argument('--smax', type=int, default=20,
help="The maximum number of contexts per word.")
# Form parameters
parser.add_argument('--nmin', type=int, default=3,
help="The minimum number of characters per ngram.")
parser.add_argument('--nmax', type=int, default=5,
help="The maximum number of characters per ngram.")
parser.add_argument('--dropout', type=float, default=0,
help="The ngram dropout probability.")
parser.add_argument('--ngram_threshold', type=int, default=4,
help="The minimum number of occurrences for an ngram to get its own embedding.")
# Visdom parameters
parser.add_argument('--visdom_port', type=int, default=8098)
parser.add_argument('--visdom_server', type=str, default=None)
args = parser.parse_args(args)
vis = visdom.Visdom(port=args.visdom_port, server=args.visdom_server) if args.visdom_server else None
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
n_gpu = torch.cuda.device_count()
if args.gradient_accumulation_steps < 1:
raise ValueError("Invalid gradient_accumulation_steps: {} < 1".format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if os.path.exists(args.output_dir) and os.listdir(args.output_dir):
raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
# generate output directories for all save_epochs
for epoch in range(int(args.num_train_epochs)):
if (epoch + 1) in args.save_epochs:
out_dir = args.output_dir + '-e' + str(epoch + 1)
if os.path.exists(out_dir) and os.listdir(out_dir):
raise ValueError("Output directory ({}) already exists and is not empty.".format(out_dir))
if not os.path.exists(out_dir):
os.makedirs(out_dir)
train_files = [args.train_dir + 'train.bucket' + str(i) + '.txt' for i in range(args.num_buckets)]
_, _, bertram_cls = bertram.MODELS[args.model_cls]
if args.bert_model in ['bert-base-uncased', 'bert-large-uncased', 'roberta-base', 'roberta-large']:
logger.info("Initializing new BERTRAM instance from {}.".format(args.bert_model))
input_processor = InputProcessor(
word_embeddings_file=args.emb_file,
word_embeddings_format=args.emb_format,
train_files=train_files,
vocab_file=args.vocab,
vector_size=args.emb_dim,
nmin=args.nmin,
nmax=args.nmax,
ngram_dropout=args.dropout,
ngram_threshold=args.ngram_threshold,
smin=args.smin,
smax=args.smax,
min_word_count=args.min_word_count,
max_seq_length=args.max_seq_length,
form_only=(args.mode == bertram.MODE_FORM),
model_cls=args.model_cls,
bert_model=args.bert_model
)
cache_dir = args.cache_dir if args.cache_dir else os.path.join(PYTORCH_TRANSFORMERS_CACHE, 'distributed_-1')
bertram_config = bertram.BertramConfig(
transformer_cls=args.model_cls, output_size=args.emb_dim, mode=args.mode,
ngram_vocab_size=input_processor.ngram_builder.get_number_of_ngrams()
)
model = bertram_cls.from_pretrained(args.bert_model, cache_dir=cache_dir, bertram_config=bertram_config)
else:
logger.info("Initializing pretrained BERTRAM instance from {}.".format(args.bert_model))
input_processor = InputProcessor.load(os.path.join(args.bert_model, IP_NAME))
bertram_config = bertram.BertramConfig.load(os.path.join(args.bert_model, CONFIG_NAME))
model, loading_info = bertram_cls.from_pretrained(args.bert_model, bertram_config=bertram_config,
output_loading_info=True) # type: Bertram
if loading_info['missing_keys']:
raise ValueError('Something went wrong loading a pretrained model: {}'.format(loading_info))
if args.mode != model.bertram_config.mode:
logger.warning("Overwriting original mode {} with {}.".format(model.bertram_config.mode, args.mode))
model.bertram_config.mode = args.mode
input_processor.mode = args.mode
model.setup()
model = model.to(device)
model_orig = model
if n_gpu > 1:
model = torch.nn.DataParallel(model)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
if args.no_finetuning:
optimizer_grouped_parameters = [{
'params': [p for n, p in param_optimizer if ('bert.' not in n and 'roberta.' not in n)],
'weight_decay': 0.01
}]
model_orig.transformer.eval()
for name, param in model_orig.transformer.named_parameters():
param.requires_grad = False
elif args.optimize_only_combinator:
optimizer_grouped_parameters = [{
'params': [p for n, p in param_optimizer if
('bert.' not in n and 'roberta.' not in n) and 'ngram_processor.' not in n],
'weight_decay': 0.01
}]
for name, param in model_orig.named_parameters():
if 'bert.' in name or 'roberta.' in name or 'ngram_processor.' in name:
param.requires_grad = False
else:
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
num_train_examples = input_processor.get_number_of_train_examples_per_epoch()
avg_contexts_per_word = (args.smin + args.smax) / 2
avg_examples_per_batch = args.train_batch_size / avg_contexts_per_word
num_train_optimization_steps = int(
num_train_examples / avg_examples_per_batch / args.gradient_accumulation_steps) * args.num_train_epochs
if args.mode == bertram.MODE_FORM:
optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
else:
optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args.warmup_proportion * num_train_optimization_steps,
t_total=num_train_optimization_steps)
global_step = 0
if vis is not None:
loss_window = vis.line(Y=np.array([0]), X=np.array([0]),
opts=dict(xlabel='Step', ylabel='Loss', title='Training loss (' + args.output_dir + ')',
legend=['Loss']))
else:
loss_window = None
model.train()
for epoch in trange(int(args.num_train_epochs), desc="Epoch"):
summed_loss = 0
step = 0
while True:
try:
nr_of_parallel_batches = max(1, n_gpu)
batches = []
# fix a number of chunks for the current set of batches to make parallel processing feasible
num_chunks = -1
for i in range(nr_of_parallel_batches):
batch = input_processor.generate_batch_from_buffer(args.train_batch_size, num_chunks=num_chunks)
batches.append(batch)
num_chunks = len(batch.nrs_of_contexts)
input_ids = torch.stack([batch.input_ids for batch in batches]).to(device)
segment_ids = torch.stack([batch.segment_ids for batch in batches]).to(device)
nrs_of_contexts = torch.stack([batch.nrs_of_contexts for batch in batches]).to(device)
mask_positions = torch.stack([batch.mask_positions for batch in batches]).to(device)
input_mask = torch.stack([batch.input_mask for batch in batches]).to(device)
# pad ngram ids
if len(batches) > 1:
max_nr_of_ngrams = max([batch.ngram_features.ngram_ids.shape[1] for batch in batches])
for batch in batches:
batch_nr_of_ngrams = batch.ngram_features.ngram_ids.shape[1]
nr_of_words = batch.ngram_features.ngram_ids.shape[0]
padding = torch.zeros((nr_of_words, max_nr_of_ngrams - batch_nr_of_ngrams), dtype=torch.long)
batch.ngram_features.ngram_ids = torch.cat([batch.ngram_features.ngram_ids, padding], dim=1)
ngram_ids = torch.stack([batch.ngram_features.ngram_ids for batch in batches]).to(device)
ngram_lengths = torch.stack([batch.ngram_features.ngram_lengths for batch in batches]).to(device)
target_vectors = torch.stack([batch.target_vectors for batch in batches]).to(device)
loss = model(input_ids, segment_ids, nrs_of_contexts, mask_positions, input_mask,
ngram_ids, ngram_lengths, target_vectors)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
loss.backward()
summed_loss += loss.item()
logger.debug('Done with backward step')
if step > 0 and step % 100 == 0:
if vis is not None:
vis.line(Y=np.array([summed_loss / 100]), X=np.array([global_step]), win=loss_window,
update='append')
logger.info('Step: %d\tLoss: %.17f', step, (summed_loss / 100))
summed_loss = 0
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
optimizer.zero_grad()
global_step += 1
step += 1
except EndOfDatasetException:
logger.info('Done with epoch %d', epoch)
input_processor.reset()
if (epoch + 1) in args.save_epochs:
out_dir = args.output_dir + '-e' + str(epoch + 1)
model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(out_dir, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(out_dir, TRANSFORMER_CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
output_ip_file = os.path.join(out_dir, IP_NAME)
input_processor.save(output_ip_file)
output_bc_file = os.path.join(out_dir, CONFIG_NAME)
model_to_save.bertram_config.save(output_bc_file)
break
# Save a trained model and the associated configuration
model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(args.output_dir, TRANSFORMER_CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
output_ip_file = os.path.join(args.output_dir, IP_NAME)
input_processor.save(output_ip_file)
output_bc_file = os.path.join(args.output_dir, CONFIG_NAME)
model_to_save.bertram_config.save(output_bc_file)