def data_loader(self, docs: Sequence[Document], batch_size, shuffle=False, label=True, **kwargs) -> DataLoader: if label is True and self.label_map is None: raise ValueError('Please specify label_map') batchify_fn = kwargs.get('batchify_fn', sequence_batchify_fn) bucket = kwargs.get('bucket', False) num_buckets = kwargs.get('num_buckets', 10) ratio = kwargs.get('ratio', 0) dataset = SequencesDataset(docs=docs, embs=self.embs, key=self.key, label_map=self.label_map, label=label) if bucket is True: dataset_lengths = list(map(lambda x: float(len(x[3])), dataset)) batch_sampler = FixedBucketSampler(dataset_lengths, batch_size=batch_size, num_buckets=num_buckets, ratio=ratio, shuffle=shuffle) return DataLoader(dataset=dataset, batch_sampler=batch_sampler, batchify_fn=batchify_fn) else: return DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle, batchify_fn=batchify_fn)
def test_sharded_data_loader_record_file(): # test record file url_format = 'https://apache-mxnet.s3-accelerate.amazonaws.com/gluon/dataset/pikachu/{}' filename = 'val.rec' idx_filename = 'val.idx' download(url_format.format(filename), path=os.path.join('tests', 'data', filename)) download(url_format.format(idx_filename), path=os.path.join('tests', 'data', idx_filename)) rec_dataset = gluon.data.vision.ImageRecordDataset( os.path.join('tests', 'data', filename)) num_workers = 2 num_shards = 4 X = np.random.uniform(size=(100, 20)) Y = np.random.uniform(size=(100, )) batch_sampler = FixedBucketSampler(lengths=[X.shape[1]] * X.shape[0], batch_size=2, num_buckets=1, shuffle=False, num_shards=num_shards) loader = ShardedDataLoader(rec_dataset, batch_sampler=batch_sampler, num_workers=num_workers) for i, seqs in enumerate(loader): assert len(seqs) == num_shards
def test_dataloader(): batch_size = 128 dataset = NLPCCDataset('train', data_root_dir) transformer = DataTransformer(['train']) word_vocab = transformer._word_vocab transformed_dataset = dataset.transform(transformer, lazy=False) batchify_fn = Tuple( Stack(), Pad(axis=0, pad_val=word_vocab[word_vocab.padding_token], ret_length=True), Stack()) sampler = FixedBucketSampler( lengths=[len(item[1]) for item in transformed_dataset], batch_size=batch_size, shuffle=True, num_buckets=30) data_loader = DataLoader(transformed_dataset, batchify_fn=batchify_fn, batch_sampler=sampler) for i, (rec_id, (data, original_length), label) in enumerate(data_loader): print(data.shape) assert data.shape[0] <= 128
def transform(raw_data, params): # 定义数据转换接口 # raw_data --> batch_data for data in raw_data: pass num_buckets = params.num_buckets batch_size = params.batch_size responses = raw_data batch_idxes = FixedBucketSampler([len(rs) for rs in responses], batch_size, num_buckets=num_buckets) batch = [] def index(r): correct = 0 if r[1] <= 0 else 1 return r[0] * 2 + correct for batch_idx in tqdm(batch_idxes, "batchify"): batch_rs = [] batch_pick_index = [] batch_labels = [] for idx in batch_idx: batch_rs.append([index(r) for r in responses[idx]]) if len(responses[idx]) <= 1: pick_index, labels = [], [] else: pick_index, labels = zip(*[(r[0], 0 if r[1] <= 0 else 1) for r in responses[idx][1:]]) batch_pick_index.append(list(pick_index)) batch_labels.append(list(labels)) max_len = max([len(rs) for rs in batch_rs]) padder = PadSequence(max_len, pad_val=0) batch_rs, data_mask = zip(*[(padder(rs), len(rs)) for rs in batch_rs]) max_len = max([len(rs) for rs in batch_labels]) padder = PadSequence(max_len, pad_val=0) batch_labels, label_mask = zip(*[(padder(labels), len(labels)) for labels in batch_labels]) batch_pick_index = [ padder(pick_index) for pick_index in batch_pick_index ] batch.append([ mx.nd.array(batch_rs), mx.nd.array(data_mask), mx.nd.array(batch_labels), mx.nd.array(batch_pick_index), mx.nd.array(label_mask) ]) return batch
def transform(raw_data, params): # 定义数据转换接口 # raw_data --> batch_data num_buckets = params.num_buckets batch_size = params.batch_size responses = raw_data # 不同长度的样本每次都会分配到固定的bucket当中 batch_idxes = FixedBucketSampler([len(rs) for rs in responses], batch_size, num_buckets=num_buckets) batch = [] def index(r): correct = 0 if r[1] <= 0 else 1 return r[0] * 2 + correct for batch_idx in tqdm(batch_idxes, "batchify"): batch_rs = [] batch_pick_index = [] batch_labels = [] for idx in batch_idx: # 1. 答对答错存储的id不一样 # 2. batch_pick_index[i], batch_labels[i]维度比batch_rs[i]小1 batch_rs.append([index(r) for r in responses[idx]]) if len(responses[idx]) <= 1: pick_index, labels = [], [] else: pick_index, labels = zip(*[(r[0], 0 if r[1] <= 0 else 1) for r in responses[idx][1:]]) batch_pick_index.append(list(pick_index)) batch_labels.append(list(labels)) # 每一个bucket中的max_len不一样 max_len = max([len(rs) for rs in batch_rs]) padder = PadSequence(max_len, pad_val=0) # batch_rs padding到max_len之后的数据 # data_mask表示对应的有效数据的条数 batch_rs, data_mask = zip(*[(padder(rs), len(rs)) for rs in batch_rs]) max_len = max([len(rs) for rs in batch_labels]) padder = PadSequence(max_len, pad_val=0) batch_labels, label_mask = zip(*[(padder(labels), len(labels)) for labels in batch_labels]) batch_pick_index = [padder(pick_index) for pick_index in batch_pick_index] # Load # 所有数据都padding到固定长度的序列 # data_mask label_mask表示有效数据的长度 # len(batch_labels) + 1 = len(batch_rs) batch.append( [torch.tensor(batch_rs), torch.tensor(data_mask), torch.tensor(batch_labels), torch.tensor(batch_pick_index), torch.tensor(label_mask)]) return batch
def get_dataloader(dataset): """create data loader based on the dataset chunk""" t0 = time.time() lengths = dataset.get_field('valid_lengths') logging.debug('Num samples = %d', len(lengths)) # A batch includes: input_id, masked_id, masked_position, masked_weight, # next_sentence_label, segment_id, valid_length batchify_fn = Tuple(Pad(), Pad(), Pad(), Pad(), Stack(), Pad(), Stack()) if args.by_token: # sharded data loader sampler = nlp.data.FixedBucketSampler( lengths=lengths, # batch_size per shard batch_size=batch_size, num_buckets=args.num_buckets, shuffle=is_train, use_average_length=True, num_shards=num_ctxes) dataloader = nlp.data.ShardedDataLoader(dataset, batch_sampler=sampler, batchify_fn=batchify_fn, num_workers=num_ctxes) logging.debug('Batch Sampler:\n%s', sampler.stats()) else: sampler = FixedBucketSampler(lengths, batch_size=batch_size * num_ctxes, num_buckets=args.num_buckets, ratio=0, shuffle=is_train) dataloader = DataLoader(dataset=dataset, batch_sampler=sampler, batchify_fn=batchify_fn, num_workers=1) logging.debug('Batch Sampler:\n%s', sampler.stats()) t1 = time.time() logging.debug('Dataloader creation cost = %.2f s', t1 - t0) return dataloader
def transform(raw_data, params): # 定义数据转换接口 # raw_data --> batch_data num_buckets = params.num_buckets batch_size = params.batch_size responses = raw_data batch_idxes = FixedBucketSampler([len(rs) for rs in responses], batch_size, num_buckets=num_buckets) batch = [] def response_index(r): correct = 0 if r[1] <= 0 else 1 return r[0] * 2 + correct def question_index(r): return r[0] for batch_idx in tqdm(batch_idxes, "batchify"): batch_qs = [] batch_rs = [] batch_labels = [] for idx in batch_idx: batch_qs.append([question_index(r) for r in responses[idx]]) batch_rs.append([response_index(r) for r in responses[idx]]) labels = [0 if r[1] <= 0 else 1 for r in responses[idx][:]] batch_labels.append(list(labels)) max_len = max([len(rs) for rs in batch_rs]) padder = PadSequence(max_len, pad_val=0) batch_qs, _ = zip(*[(padder(qs), len(qs)) for qs in batch_qs]) batch_rs, data_mask = zip(*[(padder(rs), len(rs)) for rs in batch_rs]) max_len = max([len(rs) for rs in batch_labels]) padder = PadSequence(max_len, pad_val=0) batch_labels, label_mask = zip(*[(padder(labels), len(labels)) for labels in batch_labels]) batch.append([ mx.nd.array(batch_qs, dtype="float32"), mx.nd.array(batch_rs, dtype="float32"), mx.nd.array(data_mask), mx.nd.array(batch_labels), mx.nd.array(label_mask) ]) return batch
def test_sharded_data_loader(): X = np.random.uniform(size=(100, 20)) Y = np.random.uniform(size=(100, )) dataset = gluon.data.ArrayDataset(X, Y) loader = ShardedDataLoader(dataset, 2) for i, (x, y) in enumerate(loader): assert mx.test_utils.almost_equal(x.asnumpy(), X[i * 2:(i + 1) * 2]) assert mx.test_utils.almost_equal(y.asnumpy(), Y[i * 2:(i + 1) * 2]) num_shards = 4 batch_sampler = FixedBucketSampler(lengths=[X.shape[1]] * X.shape[0], batch_size=2, num_buckets=1, shuffle=False, num_shards=num_shards) for thread_pool in [True, False]: for num_workers in [0, 1, 2, 3, 4]: loader = ShardedDataLoader(dataset, batch_sampler=batch_sampler, num_workers=num_workers, thread_pool=thread_pool) for i, seqs in enumerate(loader): assert len(seqs) == num_shards for j in range(num_shards): if i != len(loader) - 1: assert mx.test_utils.almost_equal( seqs[j][0].asnumpy(), X[(i * num_shards + j) * 2:(i * num_shards + j + 1) * 2]) assert mx.test_utils.almost_equal( seqs[j][1].asnumpy(), Y[(i * num_shards + j) * 2:(i * num_shards + j + 1) * 2]) else: assert mx.test_utils.almost_equal( seqs[j][0].asnumpy(), X[(i * num_shards + j) * 2 - num_shards:(i * num_shards + j + 1) * 2 - num_shards]) assert mx.test_utils.almost_equal( seqs[j][1].asnumpy(), Y[(i * num_shards + j) * 2 - num_shards:(i * num_shards + j + 1) * 2 - num_shards])
def transform_segment(transformer, segment, options): dataset = NLPCCDataset(segment, './data/') transformed_dataset = dataset.transform(transformer, lazy=False) word_vocab = transformer.get_word_vocab() batchify_fn = Tuple( Stack(), Pad(axis=0, pad_val=word_vocab[word_vocab.padding_token], ret_length=True), Stack()) sampler = FixedBucketSampler( lengths=[len(item[1]) for item in transformed_dataset], batch_size=options.batch_size, shuffle=True, num_buckets=options.num_buckets) return DataLoader(transformed_dataset, batchify_fn=batchify_fn, batch_sampler=sampler)
def transform(raw_data, params): # 定义数据转换接口 # raw_data --> batch_data batch_size = params.batch_size padding = params.padding num_buckets = params.num_buckets fixed_length = params.fixed_length features, labels = raw_data word_feature, word_radical_feature, char_feature, char_radical_feature = features batch_idxes = FixedBucketSampler([len(word_f) for word_f in word_feature], batch_size, num_buckets=num_buckets) batch = [] for batch_idx in batch_idxes: batch_features = [[] for _ in range(len(features))] batch_labels = [] for idx in batch_idx: for i, feature in enumerate(batch_features): batch_features[i].append(features[i][idx]) batch_labels.append(labels[idx]) batch_data = [] word_mask = [] char_mask = [] for i, feature in enumerate(batch_features): max_len = max([len(fea) for fea in feature ]) if not fixed_length else fixed_length padder = PadSequence(max_len, pad_val=padding) feature, mask = zip(*[(padder(fea), len(fea)) for fea in feature]) if i == 0: word_mask = mask elif i == 2: char_mask = mask batch_data.append(mx.nd.array(feature)) batch_data.append(mx.nd.array(word_mask)) batch_data.append(mx.nd.array(char_mask)) batch_data.append(mx.nd.array(batch_labels, dtype=np.int)) batch.append(batch_data) return batch[::-1]
def train(): """Training function.""" trainer = gluon.Trainer(model.collect_params(), args.optimizer, { 'learning_rate': args.lr, 'beta2': 0.98, 'epsilon': 1e-9 }) train_batchify_fn = btf.Tuple(btf.Pad(), btf.Pad(), btf.Stack(dtype='float32'), btf.Stack(dtype='float32')) test_batchify_fn = btf.Tuple(btf.Pad(), btf.Pad(), btf.Stack(dtype='float32'), btf.Stack(dtype='float32'), btf.Stack()) target_val_lengths = list(map(lambda x: x[-1], data_val_lengths)) target_test_lengths = list(map(lambda x: x[-1], data_test_lengths)) if args.bucket_scheme == 'constant': bucket_scheme = ConstWidthBucket() elif args.bucket_scheme == 'linear': bucket_scheme = LinearWidthBucket() elif args.bucket_scheme == 'exp': bucket_scheme = ExpWidthBucket(bucket_len_step=1.2) else: raise NotImplementedError train_batch_sampler = FixedBucketSampler(lengths=data_train_lengths, batch_size=args.batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=True, use_average_length=True, num_shards=len(ctx), bucket_scheme=bucket_scheme) logging.info('Train Batch Sampler:\n{}'.format( train_batch_sampler.stats())) train_data_loader = ShardedDataLoader(data_train, batch_sampler=train_batch_sampler, batchify_fn=train_batchify_fn, num_workers=8) val_batch_sampler = FixedBucketSampler(lengths=target_val_lengths, batch_size=args.test_batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=False, use_average_length=True, bucket_scheme=bucket_scheme) logging.info('Valid Batch Sampler:\n{}'.format(val_batch_sampler.stats())) val_data_loader = DataLoader(data_val, batch_sampler=val_batch_sampler, batchify_fn=test_batchify_fn, num_workers=8) test_batch_sampler = FixedBucketSampler(lengths=target_test_lengths, batch_size=args.test_batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=False, use_average_length=True, bucket_scheme=bucket_scheme) logging.info('Test Batch Sampler:\n{}'.format(test_batch_sampler.stats())) test_data_loader = DataLoader(data_test, batch_sampler=test_batch_sampler, batchify_fn=test_batchify_fn, num_workers=8) if args.bleu == 'tweaked': bpe = bool(args.dataset != 'IWSLT2015' and args.dataset != 'TOY') split_compound_word = bpe tokenized = True elif args.bleu == '13a' or args.bleu == 'intl': bpe = False split_compound_word = False tokenized = False else: raise NotImplementedError best_valid_bleu = 0.0 step_num = 0 warmup_steps = args.warmup_steps grad_interval = args.num_accumulated model.collect_params().setattr('grad_req', 'add') average_start = (len(train_data_loader) // grad_interval) * (args.epochs - args.average_start) average_param_dict = None model.collect_params().zero_grad() for epoch_id in range(args.epochs): log_avg_loss = 0 log_wc = 0 loss_denom = 0 step_loss = 0 log_start_time = time.time() for batch_id, seqs \ in enumerate(train_data_loader): if batch_id % grad_interval == 0: step_num += 1 new_lr = args.lr / math.sqrt(args.num_units) \ * min(1. / math.sqrt(step_num), step_num * warmup_steps ** (-1.5)) trainer.set_learning_rate(new_lr) src_wc, tgt_wc, bs = np.sum( [(shard[2].sum(), shard[3].sum(), shard[0].shape[0]) for shard in seqs], axis=0) src_wc = src_wc.asscalar() tgt_wc = tgt_wc.asscalar() loss_denom += tgt_wc - bs seqs = [[seq.as_in_context(context) for seq in shard] for context, shard in zip(ctx, seqs)] Ls = [] with mx.autograd.record(): for src_seq, tgt_seq, src_valid_length, tgt_valid_length in seqs: out, _ = model(src_seq, tgt_seq[:, :-1], src_valid_length, tgt_valid_length - 1) smoothed_label = label_smoothing(tgt_seq[:, 1:]) ls = loss_function(out, smoothed_label, tgt_valid_length - 1).sum() Ls.append((ls * (tgt_seq.shape[1] - 1)) / args.batch_size / 100.0) for L in Ls: L.backward() if batch_id % grad_interval == grad_interval - 1 or\ batch_id == len(train_data_loader) - 1: if average_param_dict is None: average_param_dict = { k: v.data(ctx[0]).copy() for k, v in model.collect_params().items() } trainer.step(float(loss_denom) / args.batch_size / 100.0) param_dict = model.collect_params() param_dict.zero_grad() if step_num > average_start: alpha = 1. / max(1, step_num - average_start) for name, average_param in average_param_dict.items(): average_param[:] += alpha * ( param_dict[name].data(ctx[0]) - average_param) step_loss += sum([L.asscalar() for L in Ls]) if batch_id % grad_interval == grad_interval - 1 or\ batch_id == len(train_data_loader) - 1: log_avg_loss += step_loss / loss_denom * args.batch_size * 100.0 loss_denom = 0 step_loss = 0 log_wc += src_wc + tgt_wc if (batch_id + 1) % (args.log_interval * grad_interval) == 0: wps = log_wc / (time.time() - log_start_time) logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, ' 'throughput={:.2f}K wps, wc={:.2f}K'.format( epoch_id, batch_id + 1, len(train_data_loader), log_avg_loss / args.log_interval, np.exp(log_avg_loss / args.log_interval), wps / 1000, log_wc / 1000)) log_start_time = time.time() log_avg_loss = 0 log_wc = 0 mx.nd.waitall() valid_loss, valid_translation_out = evaluate(val_data_loader, ctx[0]) valid_bleu_score, _, _, _, _ = compute_bleu( [val_tgt_sentences], valid_translation_out, tokenized=tokenized, tokenizer=args.bleu, split_compound_word=split_compound_word, bpe=bpe) logging.info( '[Epoch {}] valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}' .format(epoch_id, valid_loss, np.exp(valid_loss), valid_bleu_score * 100)) test_loss, test_translation_out = evaluate(test_data_loader, ctx[0]) test_bleu_score, _, _, _, _ = compute_bleu( [test_tgt_sentences], test_translation_out, tokenized=tokenized, tokenizer=args.bleu, split_compound_word=split_compound_word, bpe=bpe) logging.info( '[Epoch {}] test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'. format(epoch_id, test_loss, np.exp(test_loss), test_bleu_score * 100)) write_sentences( valid_translation_out, os.path.join(args.save_dir, 'epoch{:d}_valid_out.txt').format(epoch_id)) write_sentences( test_translation_out, os.path.join(args.save_dir, 'epoch{:d}_test_out.txt').format(epoch_id)) if valid_bleu_score > best_valid_bleu: best_valid_bleu = valid_bleu_score save_path = os.path.join(args.save_dir, 'valid_best.params') logging.info('Save best parameters to {}'.format(save_path)) model.save_parameters(save_path) save_path = os.path.join(args.save_dir, 'epoch{:d}.params'.format(epoch_id)) model.save_parameters(save_path) save_path = os.path.join(args.save_dir, 'average.params') mx.nd.save(save_path, average_param_dict) if args.average_checkpoint: for j in range(args.num_averages): params = mx.nd.load( os.path.join(args.save_dir, 'epoch{:d}.params'.format(args.epochs - j - 1))) alpha = 1. / (j + 1) for k, v in model._collect_params_with_prefix().items(): for c in ctx: v.data(c)[:] += alpha * (params[k].as_in_context(c) - v.data(c)) save_path = os.path.join( args.save_dir, 'average_checkpoint_{}.params'.format(args.num_averages)) model.save_parameters(save_path) elif args.average_start > 0: for k, v in model.collect_params().items(): v.set_data(average_param_dict[k]) save_path = os.path.join(args.save_dir, 'average.params') model.save_parameters(save_path) else: model.load_parameters(os.path.join(args.save_dir, 'valid_best.params'), ctx) valid_loss, valid_translation_out = evaluate(val_data_loader, ctx[0]) valid_bleu_score, _, _, _, _ = compute_bleu( [val_tgt_sentences], valid_translation_out, tokenized=tokenized, tokenizer=args.bleu, bpe=bpe, split_compound_word=split_compound_word) logging.info( 'Best model valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}'. format(valid_loss, np.exp(valid_loss), valid_bleu_score * 100)) test_loss, test_translation_out = evaluate(test_data_loader, ctx[0]) test_bleu_score, _, _, _, _ = compute_bleu( [test_tgt_sentences], test_translation_out, tokenized=tokenized, tokenizer=args.bleu, bpe=bpe, split_compound_word=split_compound_word) logging.info( 'Best model test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'. format(test_loss, np.exp(test_loss), test_bleu_score * 100)) write_sentences(valid_translation_out, os.path.join(args.save_dir, 'best_valid_out.txt')) write_sentences(test_translation_out, os.path.join(args.save_dir, 'best_test_out.txt'))
def train(): """Training function.""" trainer = gluon.Trainer(model.collect_params(), args.optimizer, { 'learning_rate': args.lr, 'beta2': 0.98, 'epsilon': 1e-9 }) train_batchify_fn = btf.Tuple(btf.Pad(), btf.Pad(), btf.Stack(), btf.Stack()) test_batchify_fn = btf.Tuple(btf.Pad(), btf.Pad(), btf.Stack(), btf.Stack(), btf.Stack()) target_val_lengths = list(map(lambda x: x[-1], data_val_lengths)) target_test_lengths = list(map(lambda x: x[-1], data_test_lengths)) train_batch_sampler = FixedBucketSampler(lengths=data_train_lengths, batch_size=args.batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=True, use_average_length=True) logging.info('Train Batch Sampler:\n{}'.format( train_batch_sampler.stats())) train_data_loader = DataLoader(data_train, batch_sampler=train_batch_sampler, batchify_fn=train_batchify_fn, num_workers=8) val_batch_sampler = FixedBucketSampler(lengths=target_val_lengths, batch_size=args.test_batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=False, use_average_length=True) logging.info('Valid Batch Sampler:\n{}'.format(val_batch_sampler.stats())) val_data_loader = DataLoader(data_val, batch_sampler=val_batch_sampler, batchify_fn=test_batchify_fn, num_workers=8) test_batch_sampler = FixedBucketSampler(lengths=target_test_lengths, batch_size=args.test_batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=False, use_average_length=True) logging.info('Test Batch Sampler:\n{}'.format(test_batch_sampler.stats())) test_data_loader = DataLoader(data_test, batch_sampler=test_batch_sampler, batchify_fn=test_batchify_fn, num_workers=8) best_valid_bleu = 0.0 step_num = 0 warmup_steps = args.warmup_steps grad_interval = args.num_accumulated model.collect_params().setattr('grad_req', 'add') average_start = (len(train_data_loader) // grad_interval) * (args.epochs - args.average_start) average_param_dict = None model.collect_params().zero_grad() for epoch_id in range(args.epochs): log_avg_loss = 0 log_wc = 0 loss_denom = 0 step_loss = 0 log_start_time = time.time() for batch_id, (src_seq, tgt_seq, src_valid_length, tgt_valid_length) \ in enumerate(train_data_loader): if batch_id % grad_interval == 0: step_num += 1 new_lr = args.lr / math.sqrt(args.num_units) \ * min(1. / math.sqrt(step_num), step_num * warmup_steps ** (-1.5)) trainer.set_learning_rate(new_lr) # logging.info(src_seq.context) Context suddenly becomes GPU. src_wc = src_valid_length.sum().asscalar() tgt_wc = tgt_valid_length.sum().asscalar() loss_denom += tgt_wc - tgt_valid_length.shape[0] if src_seq.shape[0] > len(ctx): src_seq_list, tgt_seq_list, src_valid_length_list, tgt_valid_length_list \ = [gluon.utils.split_and_load(seq, ctx, batch_axis=0, even_split=False) for seq in [src_seq, tgt_seq, src_valid_length, tgt_valid_length]] else: src_seq_list = [src_seq.as_in_context(ctx[0])] tgt_seq_list = [tgt_seq.as_in_context(ctx[0])] src_valid_length_list = [ src_valid_length.as_in_context(ctx[0]) ] tgt_valid_length_list = [ tgt_valid_length.as_in_context(ctx[0]) ] Ls = [] with mx.autograd.record(): for src_seq, tgt_seq, src_valid_length, tgt_valid_length \ in zip(src_seq_list, tgt_seq_list, src_valid_length_list, tgt_valid_length_list): out, _ = model(src_seq, tgt_seq[:, :-1], src_valid_length, tgt_valid_length - 1) smoothed_label = label_smoothing(tgt_seq[:, 1:]) ls = loss_function(out, smoothed_label, tgt_valid_length - 1).sum() Ls.append((ls * (tgt_seq.shape[1] - 1)) / args.batch_size) for L in Ls: L.backward() if batch_id % grad_interval == grad_interval - 1 or\ batch_id == len(train_data_loader) - 1: if average_param_dict is None: average_param_dict = { k: v.data(ctx[0]).copy() for k, v in model.collect_params().items() } trainer.step(float(loss_denom) / args.batch_size) param_dict = model.collect_params() param_dict.zero_grad() if step_num > average_start: alpha = 1. / max(1, step_num - average_start) for name, average_param in average_param_dict.items(): average_param[:] += alpha * ( param_dict[name].data(ctx[0]) - average_param) step_loss += sum([L.asscalar() for L in Ls]) if batch_id % grad_interval == grad_interval - 1 or\ batch_id == len(train_data_loader) - 1: log_avg_loss += step_loss / loss_denom * args.batch_size loss_denom = 0 step_loss = 0 log_wc += src_wc + tgt_wc if (batch_id + 1) % (args.log_interval * grad_interval) == 0: wps = log_wc / (time.time() - log_start_time) logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, ' 'throughput={:.2f}K wps, wc={:.2f}K'.format( epoch_id, batch_id + 1, len(train_data_loader), log_avg_loss / args.log_interval, np.exp(log_avg_loss / args.log_interval), wps / 1000, log_wc / 1000)) log_start_time = time.time() log_avg_loss = 0 log_wc = 0 mx.nd.waitall() valid_loss, valid_translation_out = evaluate(val_data_loader, ctx[0]) valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out, bpe=True, split_compound_word=True) logging.info( '[Epoch {}] valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}' .format(epoch_id, valid_loss, np.exp(valid_loss), valid_bleu_score * 100)) test_loss, test_translation_out = evaluate(test_data_loader, ctx[0]) test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out, bpe=True, split_compound_word=True) logging.info( '[Epoch {}] test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'. format(epoch_id, test_loss, np.exp(test_loss), test_bleu_score * 100)) write_sentences( valid_translation_out, os.path.join(args.save_dir, 'epoch{:d}_valid_out.txt').format(epoch_id)) write_sentences( test_translation_out, os.path.join(args.save_dir, 'epoch{:d}_test_out.txt').format(epoch_id)) if valid_bleu_score > best_valid_bleu: best_valid_bleu = valid_bleu_score save_path = os.path.join(args.save_dir, 'valid_best.params') logging.info('Save best parameters to {}'.format(save_path)) model.save_params(save_path) save_path = os.path.join(args.save_dir, 'epoch{:d}.params'.format(epoch_id)) model.save_params(save_path) save_path = os.path.join(args.save_dir, 'average.params') mx.nd.save(save_path, average_param_dict) if args.average_checkpoint: for j in range(args.num_averages): params = mx.nd.load( os.path.join(args.save_dir, 'epoch{:d}.params'.format(args.epochs - j - 1))) alpha = 1. / (j + 1) for k, v in model._collect_params_with_prefix().items(): for c in ctx: v.data(c)[:] += alpha * (params[k].as_in_context(c) - v.data(c)) elif args.average_start > 0: for k, v in model.collect_params().items(): v.set_data(average_param_dict[k]) else: model.load_params(os.path.join(args.save_dir, 'valid_best.params'), ctx) valid_loss, valid_translation_out = evaluate(val_data_loader, ctx[0]) valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out, bpe=True, split_compound_word=True) logging.info( 'Best model valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}'. format(valid_loss, np.exp(valid_loss), valid_bleu_score * 100)) test_loss, test_translation_out = evaluate(test_data_loader, ctx[0]) test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out, bpe=True, split_compound_word=True) logging.info( 'Best model test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'. format(test_loss, np.exp(test_loss), test_bleu_score * 100)) write_sentences(valid_translation_out, os.path.join(args.save_dir, 'best_valid_out.txt')) write_sentences(test_translation_out, os.path.join(args.save_dir, 'best_test_out.txt'))
def train(): """Training function.""" trainer = gluon.Trainer(model.collect_params(), args.optimizer, {'learning_rate': args.lr}) train_batchify_fn = btf.Tuple(btf.Pad(), btf.Pad(), btf.Stack(dtype='float32'), btf.Stack(dtype='float32')) test_batchify_fn = btf.Tuple(btf.Pad(), btf.Pad(), btf.Stack(dtype='float32'), btf.Stack(dtype='float32'), btf.Stack()) if args.bucket_scheme == 'constant': bucket_scheme = ConstWidthBucket() elif args.bucket_scheme == 'linear': bucket_scheme = LinearWidthBucket() elif args.bucket_scheme == 'exp': bucket_scheme = ExpWidthBucket(bucket_len_step=1.2) else: raise NotImplementedError train_batch_sampler = FixedBucketSampler(lengths=data_train_lengths, batch_size=args.batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=True, bucket_scheme=bucket_scheme) logging.info('Train Batch Sampler:\n{}'.format( train_batch_sampler.stats())) train_data_loader = DataLoader(data_train, batch_sampler=train_batch_sampler, batchify_fn=train_batchify_fn, num_workers=8) val_batch_sampler = FixedBucketSampler(lengths=data_val_lengths, batch_size=args.test_batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=False) logging.info('Valid Batch Sampler:\n{}'.format(val_batch_sampler.stats())) val_data_loader = DataLoader(data_val, batch_sampler=val_batch_sampler, batchify_fn=test_batchify_fn, num_workers=8) test_batch_sampler = FixedBucketSampler(lengths=data_test_lengths, batch_size=args.test_batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=False) logging.info('Test Batch Sampler:\n{}'.format(test_batch_sampler.stats())) test_data_loader = DataLoader(data_test, batch_sampler=test_batch_sampler, batchify_fn=test_batchify_fn, num_workers=8) best_valid_bleu = 0.0 for epoch_id in range(args.epochs): log_avg_loss = 0 log_avg_gnorm = 0 log_wc = 0 log_start_time = time.time() for batch_id, (src_seq, tgt_seq, src_valid_length, tgt_valid_length)\ in enumerate(train_data_loader): # logging.info(src_seq.context) Context suddenly becomes GPU. src_seq = src_seq.as_in_context(ctx) tgt_seq = tgt_seq.as_in_context(ctx) src_valid_length = src_valid_length.as_in_context(ctx) tgt_valid_length = tgt_valid_length.as_in_context(ctx) with mx.autograd.record(): out, _ = model(src_seq, tgt_seq[:, :-1], src_valid_length, tgt_valid_length - 1) loss = loss_function(out, tgt_seq[:, 1:], tgt_valid_length - 1).mean() loss = loss * (tgt_seq.shape[1] - 1) / (tgt_valid_length - 1).mean() loss.backward() grads = [p.grad(ctx) for p in model.collect_params().values()] gnorm = gluon.utils.clip_global_norm(grads, args.clip) trainer.step(1) src_wc = src_valid_length.sum().asscalar() tgt_wc = (tgt_valid_length - 1).sum().asscalar() step_loss = loss.asscalar() log_avg_loss += step_loss log_avg_gnorm += gnorm log_wc += src_wc + tgt_wc if (batch_id + 1) % args.log_interval == 0: wps = log_wc / (time.time() - log_start_time) logging.info( '[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, gnorm={:.4f}, ' 'throughput={:.2f}K wps, wc={:.2f}K'.format( epoch_id, batch_id + 1, len(train_data_loader), log_avg_loss / args.log_interval, np.exp(log_avg_loss / args.log_interval), log_avg_gnorm / args.log_interval, wps / 1000, log_wc / 1000)) log_start_time = time.time() log_avg_loss = 0 log_avg_gnorm = 0 log_wc = 0 valid_loss, valid_translation_out = evaluate(val_data_loader) valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out) logging.info( '[Epoch {}] valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}' .format(epoch_id, valid_loss, np.exp(valid_loss), valid_bleu_score * 100)) test_loss, test_translation_out = evaluate(test_data_loader) test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out) logging.info( '[Epoch {}] test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'. format(epoch_id, test_loss, np.exp(test_loss), test_bleu_score * 100)) write_sentences( valid_translation_out, os.path.join(args.save_dir, 'epoch{:d}_valid_out.txt').format(epoch_id)) write_sentences( test_translation_out, os.path.join(args.save_dir, 'epoch{:d}_test_out.txt').format(epoch_id)) if valid_bleu_score > best_valid_bleu: best_valid_bleu = valid_bleu_score save_path = os.path.join(args.save_dir, 'valid_best.params') logging.info('Save best parameters to {}'.format(save_path)) model.save_params(save_path) if epoch_id + 1 >= (args.epochs * 2) // 3: new_lr = trainer.learning_rate * args.lr_update_factor logging.info('Learning rate change to {}'.format(new_lr)) trainer.set_learning_rate(new_lr) model.load_params(os.path.join(args.save_dir, 'valid_best.params')) valid_loss, valid_translation_out = evaluate(val_data_loader) valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out) logging.info( 'Best model valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}'. format(valid_loss, np.exp(valid_loss), valid_bleu_score * 100)) test_loss, test_translation_out = evaluate(test_data_loader) test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out) logging.info( 'Best model test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'. format(test_loss, np.exp(test_loss), test_bleu_score * 100)) write_sentences(valid_translation_out, os.path.join(args.save_dir, 'best_valid_out.txt')) write_sentences(test_translation_out, os.path.join(args.save_dir, 'best_test_out.txt'))
def train(): """Training function.""" trainer = gluon.Trainer(model.collect_params(), args.optimizer, {'learning_rate': args.lr}) train_batchify_fn = btf.Tuple(btf.Pad(), btf.Pad(), btf.Stack(), btf.Stack()) test_batchify_fn = btf.Tuple(btf.Pad(), btf.Pad(), btf.Stack(), btf.Stack(), btf.Stack()) train_batch_sampler = FixedBucketSampler(lengths=data_train_lengths, batch_size=args.batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=True) logging.info('Train Batch Sampler:\n{}'.format(train_batch_sampler.stats())) train_data_loader = DataLoader(data_train, batch_sampler=train_batch_sampler, batchify_fn=train_batchify_fn, num_workers=8) val_batch_sampler = FixedBucketSampler(lengths=data_val_lengths, batch_size=args.test_batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=False) logging.info('Valid Batch Sampler:\n{}'.format(val_batch_sampler.stats())) val_data_loader = DataLoader(data_val, batch_sampler=val_batch_sampler, batchify_fn=test_batchify_fn, num_workers=8) test_batch_sampler = FixedBucketSampler(lengths=data_test_lengths, batch_size=args.test_batch_size, num_buckets=args.num_buckets, ratio=args.bucket_ratio, shuffle=False) logging.info('Test Batch Sampler:\n{}'.format(test_batch_sampler.stats())) test_data_loader = DataLoader(data_test, batch_sampler=test_batch_sampler, batchify_fn=test_batchify_fn, num_workers=8) best_valid_bleu = 0.0 for epoch_id in range(args.epochs): log_avg_loss = 0 log_avg_gnorm = 0 log_wc = 0 log_start_time = time.time() for batch_id, (src_seq, tgt_seq, src_valid_length, tgt_valid_length)\ in enumerate(train_data_loader): # logging.info(src_seq.context) Context suddenly becomes GPU. src_seq = src_seq.as_in_context(ctx) tgt_seq = tgt_seq.as_in_context(ctx) src_valid_length = src_valid_length.as_in_context(ctx) tgt_valid_length = tgt_valid_length.as_in_context(ctx) with mx.autograd.record(): out, _ = model(src_seq, tgt_seq[:, :-1], src_valid_length, tgt_valid_length - 1) loss = loss_function(out, tgt_seq[:, 1:], tgt_valid_length - 1).mean() loss = loss * (tgt_seq.shape[1] - 1) / (tgt_valid_length - 1).mean() loss.backward() grads = [p.grad(ctx) for p in model.collect_params().values()] gnorm = gluon.utils.clip_global_norm(grads, args.clip) trainer.step(1) src_wc = src_valid_length.sum().asscalar() tgt_wc = (tgt_valid_length - 1).sum().asscalar() step_loss = loss.asscalar() log_avg_loss += step_loss log_avg_gnorm += gnorm log_wc += src_wc + tgt_wc if (batch_id + 1) % args.log_interval == 0: wps = log_wc / (time.time() - log_start_time) logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, gnorm={:.4f}, ' 'throughput={:.2f}K wps, wc={:.2f}K' .format(epoch_id, batch_id + 1, len(train_data_loader), log_avg_loss / args.log_interval, np.exp(log_avg_loss / args.log_interval), log_avg_gnorm / args.log_interval, wps / 1000, log_wc / 1000)) log_start_time = time.time() log_avg_loss = 0 log_avg_gnorm = 0 log_wc = 0 valid_loss, valid_translation_out = evaluate(val_data_loader) valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out) logging.info('[Epoch {}] valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}' .format(epoch_id, valid_loss, np.exp(valid_loss), valid_bleu_score * 100)) test_loss, test_translation_out = evaluate(test_data_loader) test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out) logging.info('[Epoch {}] test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}' .format(epoch_id, test_loss, np.exp(test_loss), test_bleu_score * 100)) write_sentences(valid_translation_out, os.path.join(args.save_dir, 'epoch{:d}_valid_out.txt').format(epoch_id)) write_sentences(test_translation_out, os.path.join(args.save_dir, 'epoch{:d}_test_out.txt').format(epoch_id)) if valid_bleu_score > best_valid_bleu: best_valid_bleu = valid_bleu_score save_path = os.path.join(args.save_dir, 'valid_best.params') logging.info('Save best parameters to {}'.format(save_path)) model.save_params(save_path) else: new_lr = trainer.learning_rate * args.lr_update_factor logging.info('Learning rate change to {}'.format(new_lr)) trainer.set_learning_rate(new_lr) model.load_params(os.path.join(args.save_dir, 'valid_best.params')) valid_loss, valid_translation_out = evaluate(val_data_loader) valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out) logging.info('Best model valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}' .format(valid_loss, np.exp(valid_loss), valid_bleu_score * 100)) test_loss, test_translation_out = evaluate(test_data_loader) test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out) logging.info('Best model test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}' .format(test_loss, np.exp(test_loss), test_bleu_score * 100)) write_sentences(valid_translation_out, os.path.join(args.save_dir, 'best_valid_out.txt')) write_sentences(test_translation_out, os.path.join(args.save_dir, 'best_test_out.txt'))