传入参数:
model:
# def parse_args 中,命令行输入模型名称,默认是"transformer"。
parser.add_argument("--model", type=str, required=True,
help="Name of the model.")
# def main 中
# model_cls = models.get_model(args.model)
返回thumt.models.transformer.Transformer类
# model = model_cls(params).cuda()
实例化Transformer类,把参数传进去初始化。
# if args.half:
model = model.half()
是否使用半精度
# model.train()
模型进入训练状态
sorted_key:
eval_dataset:
# params.validation: "newsdev2017.tc.32k.zh" 验证集
# "references": "newsdev2017.tc.en" 参考
# def main中
if params.validation:
sorted_key, eval_dataset = data.MTPipeline.get_infer_dataset(
params.validation, params)
references = load_references(params.references)
else:
sorted_key = None
eval_dataset = None
references = None
`sorted_key`变量则存储了评估数据集中的句子键值,用于将评估结果与参考翻译进行匹配。
使用`MTPipeline`类中的`get_infer_dataset`方法来创建一个评估数据集对象,并将其存储在`eval_dataset`变量中。
如果`params.validation`参数为假,则将`sorted_key`、`eval_dataset`和`references`变量都设置为`None`,表示没有评估数据集和参考翻译可用。
params.output:
# def parse_args中
parser.add_argument("--output", type=str, default="train",
help="Path to load/store checkpoints.")
输出的文件路径
references:
# "references": "newsdev2017.tc.en" 参考
# def main中
references = load_references(params.references)
def load_references(pattern):
if not pattern:
return None
files = glob.glob(pattern)
references = []
for name in files:
ref = []
with open(name, "rb") as fd:
for line in fd:
items = line.strip().split()
ref.append(items)
references.append(ref)
return list(zip(*references))
load_references
的 Python 函数,它接受一个名为 pattern
的参数。它的功能如下:
- 如果
pattern
是空字符串或 None,则返回None
。 - 它使用
glob
模块来查找与给定模式匹配的所有文件名。glob
函数返回一个文件名列表。 - 它创建一个空列表
references
。 - 它循环遍历文件名列表中的每个文件名。
- 它创建一个名为
ref
的空列表。 - 它以二进制模式打开文件,并循环遍历文件中的每一行。
- 它从行中去除换行符并将其拆分成一个项目列表。
- 它将项目列表添加到
ref
列表中。 - 它将
ref
列表添加到references
列表中。 - 它返回一个元组列表,其中每个元组包含
references
中每个列表中相应行的参考数据。使用zip
函数将列表的列表转置为元组的列表。
总体而言,该函数的设计目的是读取多个包含参考数据的文本文件,并将它们组合成一个元组列表,其中每个元组包含来自每个文件对应行的参考数据。
训练过程中:
if counter % params.update_cycle == 0: # 这段代码用于检查是否需要更新模型。
if step >= params.train_steps:
utils.evaluate(model, sorted_key, eval_dataset,
params.output, references, params)
save_checkpoint(step, epoch, model, optimizer, params)
if dist.get_rank() == 0:
summary.close()
return
'''
在这两个条件语句中,如果 step 大于等于 params.train_steps,就会执行 utils.evaluate 函数来评估模型,并且保存模型的检查点。如果当前进程是分布式训练的主进程,那么还会关闭 summary 对象。最后,整个函数返回。
'''
if step % params.eval_steps == 0: # 这段代码用于检查是否需要评估模型。
utils.evaluate(model, sorted_key, eval_dataset,
params.output, references, params)
代码解释:
# if step >= params.train_steps
如果当前步数大于等于设置的训练步数
> 调用评估函数 utils.evaluate
> 保存检查点 save_checkpoint
# if step % params.eval_steps == 0
如果当前步数取余验证集步数等于0
> 调用评估函数 utils.evaluate
评估函数:
def evaluate(model, sorted_key, dataset, base_dir, references, params):
if not references:
return
base_dir = base_dir.rstrip("/")
save_path = os.path.join(base_dir, "eval")
record_name = os.path.join(save_path, "record")
log_name = os.path.join(save_path, "log")
max_to_keep = params.keep_top_checkpoint_max
if dist.get_rank() == 0:
# Create directory and copy files
if not os.path.exists(save_path):
print("Making dir: %s" % save_path)
os.makedirs(save_path)
params_pattern = os.path.join(base_dir, "*.json")
params_files = glob.glob(params_pattern)
for name in params_files:
new_name = name.replace(base_dir, save_path)
shutil.copy(name, new_name)
# Do validation here
global_step = get_global_step()
if dist.get_rank() == 0:
print("Validating model at step %d" % global_step)
score = _evaluate_model(model, sorted_key, dataset, references, params)
# Save records
if dist.get_rank() == 0:
scalar("BLEU/score", score, global_step, write_every_n_steps=1)
print("BLEU at step %d: %f" % (global_step, score))
# Save checkpoint to save_path
save({"model": model.state_dict(), "step": global_step}, save_path)
_save_log(log_name, ("BLEU", global_step, score))
records = _read_score_record(record_name)
record = [latest_checkpoint(save_path).split("/")[-1], score]
added, removed, records = _add_to_record(records, record, max_to_keep)
if added is None:
# Remove latest checkpoint
filename = latest_checkpoint(save_path)
print("Removing %s" % filename)
files = glob.glob(filename + "*")
for name in files:
os.remove(name)
if removed is not None:
filename = os.path.join(save_path, removed)
print("Removing %s" % filename)
files = glob.glob(filename + "*")
for name in files:
os.remove(name)
_save_score_record(record_name, records)
best_score = records[0][1]
print("Best score at step %d: %f" % (global_step, best_score))
这是一个 Python 函数,名为 evaluate
。该函数接受以下参数:
model
:神经网络模型对象。sorted_key
:输入数据的排序键。dataset
:评估数据集。base_dir
:输出目录的基本路径。references
:参考数据列表。params
:参数对象。
该函数的主要目的是评估模型在给定数据集上的性能,并保存检查点、日志和得分记录。函数的大部分代码都是管理输出目录和文件的逻辑。下面是函数的主要步骤:
- 如果参考数据列表为空,则直接返回。
if not references:
return
- 确定输出目录、记录文件和日志文件的路径。
base_dir = base_dir.rstrip("/")
save_path = os.path.join(base_dir, "eval")
record_name = os.path.join(save_path, "record")
log_name = os.path.join(save_path, "log")
max_to_keep = params.keep_top_checkpoint_max
首先,base_dir
基本路径可能以斜杠结尾,因此使用 rstrip()
函数删除任何可能的结尾斜杠。
然后,将 eval
附加到 base_dir
以创建输出目录的路径,并将路径存储在 save_path
中。
record_name
和 log_name
分别是记录文件和日志文件的路径。它们都在 save_path
下的相应文件名中创建。
最后,max_to_keep
是一个整数,表示要保留的最大检查点文件数量。它从 params
参数中获取,这可能是一个包含其他参数的对象。
- 如果当前进程是分布式训练的主进程,则创建输出目录,并将参数文件复制到输出目录中。
if dist.get_rank() == 0:
# Create directory and copy files
if not os.path.exists(save_path):
print("Making dir: %s" % save_path)
os.makedirs(save_path)
params_pattern = os.path.join(base_dir, "*.json")
params_files = glob.glob(params_pattern)
for name in params_files:
new_name = name.replace(base_dir, save_path)
shutil.copy(name, new_name)
这是 evaluate
函数的一部分,用于在分布式训练中仅在主进程上创建目录并将参数文件复制到输出目录中。
首先,使用 dist.get_rank()
函数检查当前进程是否为主进程(dist.get_rank()
返回当前进程在分布式训练中的排名,如果排名为 0,则该进程是主进程)。
如果当前进程是主进程,则检查输出目录是否存在。如果不存在,则使用 os.makedirs
函数创建输出目录,并打印一条消息指示目录已创建。
接下来,使用 glob.glob
函数查找基本路径下的所有 JSON 参数文件。然后,使用 shutil.copy
函数将每个参数文件复制到输出目录中,并用新路径替换原始路径。
由于所有其他进程都不需要输出目录或参数文件,因此仅在主进程上执行此操作可以节省计算资源和时间。
- 获取全局步数
global_step
。
# Do validation here
global_step = get_global_step()
这是 evaluate
函数的一部分,用于获取全局步数 global_step
,然后对模型进行评估。
get_global_step()
函数是一个用于获取全局训练步数的函数。具体实现取决于训练框架,但通常是从训练状态中获取的。global_step
是一个表示训练步数的整数变量。
在获取全局步数之后,该函数会继续调用 _evaluate_model
函数来评估模型在给定数据集上的性能,并返回 BLEU 分数。
- 如果当前进程是分布式训练的主进程,则输出当前模型正在进行验证的信息。
if dist.get_rank() == 0:
print("Validating model at step %d" % global_step)
这是 evaluate
函数的一部分,用于在分布式训练中仅在主进程上打印消息,指示模型正在进行验证。
首先,使用 dist.get_rank()
函数检查当前进程是否为主进程(dist.get_rank()
返回当前进程在分布式训练中的排名,如果排名为 0,则该进程是主进程)。
如果当前进程是主进程,则打印一条消息,指示模型正在进行验证,并附加全局步数 global_step
信息。这可以帮助用户跟踪验证的进度和当前模型的状态。
- 调用
_evaluate_model
函数来评估模型在给定数据集上的性能,并返回 BLEU 分数。
score = _evaluate_model(model, sorted_key, dataset, references, params)
这是 evaluate
函数的一部分,用于调用 _evaluate_model
函数来评估模型在给定数据集上的性能,并返回 BLEU 分数。
_evaluate_model
函数是一个用于评估模型在给定数据集上的性能的辅助函数,它接受以下参数:
model
:神经网络模型对象。sorted_key
:输入数据的排序键。dataset
:评估数据集。references
:参考数据列表。params
:参数对象。
该函数主要目的是计算模型在给定数据集上的 BLEU 分数,并返回该分数。它使用 references
列表中的参考数据与模型生成的输出进行比较,并计算 BLEU 分数。
在 evaluate
函数中,使用 _evaluate_model
函数来计算模型在给定数据集上的 BLEU 分数,并将该分数存储在 score
变量中以备后续使用。
- 如果当前进程是分布式训练的主进程,则输出 BLEU 分数,并保存检查点、日志和得分记录。
# Save records
if dist.get_rank() == 0:
scalar("BLEU/score", score, global_step, write_every_n_steps=1)
print("BLEU at step %d: %f" % (global_step, score))
# Save checkpoint to save_path
save({"model": model.state_dict(), "step": global_step}, save_path)
_save_log(log_name, ("BLEU", global_step, score))
records = _read_score_record(record_name)
record = [latest_checkpoint(save_path).split("/")[-1], score]
added, removed, records = _add_to_record(records, record, max_to_keep)
这是 evaluate
函数的一部分,用于保存检查点、日志和得分记录。
首先,使用 dist.get_rank()
函数检查当前进程是否为主进程(dist.get_rank()
返回当前进程在分布式训练中的排名,如果排名为 0,则该进程是主进程)。
如果当前进程是主进程,则使用 scalar
函数记录 BLEU 分数,并打印一条消息,指示 BLEU 分数和全局步数 global_step
。
然后,使用 save
函数将当前模型的检查点保存到输出目录中。save
函数接受一个字典作为参数,其中包含要保存的对象以及它们的键。
接下来,使用 _save_log
函数将 BLEU 分数记录到日志文件中。_save_log
函数接受日志文件名和要记录的数据元组作为参数,并将数据元组写入日志文件中。
然后,使用 _read_score_record
函数读取得分记录文件中的数据,并使用 _add_to_record
函数将当前检查点和 BLEU 分数添加到记录中。_add_to_record
函数首先检查记录中的检查点数量是否超过了最大保留数量,如果超过了,则删除最旧的检查点文件。然后,它将新的检查点文件和 BLEU 分数添加到记录中,并返回更新的记录。
最后,added
和 removed
变量包含添加和删除的检查点文件数量。
- 如果当前进程是分布式训练的主进程,则根据得分记录和最大保留数量,删除旧的检查点文件。
if added is None:
# Remove latest checkpoint
filename = latest_checkpoint(save_path)
print("Removing %s" % filename)
files = glob.glob(filename + "*")
for name in files:
os.remove(name)
这是 evaluate
函数的一部分,用于在保留的检查点文件数量超出最大值时,从输出目录中删除最旧的检查点文件。
首先,检查 added
是否为 None
。如果是 None
,则说明没有添加新的检查点文件,因此需要删除最新的检查点文件。
然后,使用 latest_checkpoint
函数找到最新的检查点文件,并将其文件名存储在 filename
变量中。
接下来,使用 glob.glob
函数查找与文件名匹配的所有文件,并将它们的文件名存储在 files
列表中。
最后,使用 os.remove
函数删除所有文件。这将从输出目录中删除最新的检查点文件及其相关文件。
if removed is not None:
filename = os.path.join(save_path, removed)
print("Removing %s" % filename)
files = glob.glob(filename + "*")
for name in files:
os.remove(name)
这是 evaluate
函数的一部分,用于在保留的检查点文件数量超出最大值时,从输出目录中删除最旧的检查点文件。
如果 removed
不是 None
,则说明最旧的检查点文件已经被删除,因此需要删除该文件及其相关文件。
首先,使用 os.path.join
函数将输出目录和要删除的文件名合并为完整路径,并将其存储在 filename
变量中。
然后,使用 glob.glob
函数查找与文件名匹配的所有文件,并将它们的文件名存储在 files
列表中。
最后,使用 os.remove
函数删除所有文件。这将从输出目录中删除最旧的检查点文件及其相关文件。
- 保存新的检查点文件,更新得分记录,并输出最佳分数的信息。
_save_score_record(record_name, records)
best_score = records[0][1]
print("Best score at step %d: %f" % (global_step, best_score))
这是 evaluate
函数的一部分,用于保存得分记录,并打印出最佳分数和全局步数 global_step
。
首先,使用 _save_score_record
函数将更新后的得分记录保存到文件中。_save_score_record
函数接受得分记录文件名和要保存的记录列表作为参数,并将记录列表写入文件中。
然后,使用 best_score
变量存储得分记录中的最佳分数。由于得分记录已按降序排序,因此第一个记录的分数应该是最佳分数。
最后,打印一条消息,指示最佳分数和全局步数 global_step
。这可以帮助用户了解模型的最佳性能以及它发生的时间步。
评估函数的核心实现
def _evaluate_model(model, sorted_key, dataset, references, params):
# Create model
with torch.no_grad():
model.eval()
iterator = iter(dataset)
counter = 0
pad_max = 1024
# Buffers for synchronization
size = torch.zeros([dist.get_world_size()]).long()
t_list = [torch.empty([params.decode_batch_size, pad_max]).long()
for _ in range(dist.get_world_size())]
results = []
while True:
try:
features = next(iterator)
batch_size = features["source"].shape[0]
except:
features = {
"source": torch.ones([1, 1]).long(),
"source_mask": torch.ones([1, 1]).float()
}
batch_size = 0
t = time.time()
counter += 1
# Decode
seqs, _ = beam_search([model], features, params)
# Padding
seqs = torch.squeeze(seqs, dim=1)
pad_batch = params.decode_batch_size - seqs.shape[0]
pad_length = pad_max - seqs.shape[1]
seqs = torch.nn.functional.pad(seqs, (0, pad_length, 0, pad_batch))
# Synchronization
size.zero_()
size[dist.get_rank()].copy_(torch.tensor(batch_size))
dist.all_reduce(size)
dist.all_gather(t_list, seqs)
if size.sum() == 0:
break
if dist.get_rank() != 0:
continue
for i in range(params.decode_batch_size):
for j in range(dist.get_world_size()):
n = size[j]
seq = _convert_to_string(t_list[j][i], params)
if i >= n:
continue
# Restore BPE segmentation
seq = BPE.decode(seq)
results.append(seq.split())
t = time.time() - t
print("Finished batch: %d (%.3f sec)" % (counter, t))
model.train()
if dist.get_rank() == 0:
restored_results = []
for idx in range(len(results)):
restored_results.append(results[sorted_key[idx]])
return bleu(restored_results, references)
return 0.0
这是 _evaluate_model
函数,用于评估模型在给定数据集上的性能,并返回 BLEU 分数。
该函数接受以下参数:
model
:神经网络模型对象。sorted_key
:输入数据的排序键。dataset
:评估数据集。references
:参考数据列表。params
:参数对象。
该函数的主要工作流程如下:
-
将模型设置为评估模式,并使用
iter
函数创建数据集迭代器。 -
循环迭代数据集,对每个批次进行解码和填充。
-
对解码后的结果进行同步,以便在分布式训练中进行合并。
-
在主进程中,将解码结果还原为原始数据。
-
使用
bleu
函数计算还原后的结果与参考数据之间的 BLEU 分数,并返回该分数。
在 _evaluate_model
函数中,使用 beam_search
函数对输入进行解码。beam_search
函数接受神经网络模型对象、输入数据和参数对象作为参数,并使用束搜索算法生成最佳解码序列。解码后的结果是一个张量,大小为 [batch_size, seq_len]
,其中 batch_size
是批次大小,seq_len
是最大序列长度。
然后,使用 torch.squeeze
函数将解码后的结果张量的第二个维度压缩,以便将其转换为大小为 [batch_size, seq_len]
的二维张量。然后,使用 torch.nn.functional.pad
函数将解码后的结果填充为固定大小的张量,以便在分布式训练中进行合并。
最后,使用 bleu
函数计算还原后的结果与参考数据之间的 BLEU 分数,并将其返回。
# Create model
with torch.no_grad():
model.eval()
iterator = iter(dataset)
counter = 0
pad_max = 1024
这是 _evaluate_model
函数的一部分,用于创建模型,并将其设置为评估模式。
使用 torch.no_grad()
上下文管理器,禁用梯度计算,以减少内存和计算开销。
然后,使用 iter
函数创建数据集迭代器,并将迭代器存储在 iterator
变量中。
counter
变量用于记录迭代次数,而 pad_max
变量则定义了填充后的序列长度。在该代码段中,将其设置为 1024。
最后,使用 model.eval()
将模型设置为评估模式。在评估模式下,模型不会进行反向传播,这有助于减少内存和计算开销。
# Buffers for synchronization
size = torch.zeros([dist.get_world_size()]).long()
t_list = [torch.empty([params.decode_batch_size, pad_max]).long()
for _ in range(dist.get_world_size())]
results = []
这是 _evaluate_model
函数的一部分,用于设置用于同步的缓冲区。
size
变量是一个大小为 [dist.get_world_size()]
的长整型张量,用于存储每个进程的批次大小。
t_list
变量是一个大小为 [dist.get_world_size(), params.decode_batch_size, pad_max]
的三维长整型张量,用于存储每个进程的解码结果。在每个批次中,每个进程都会生成一个解码结果,因此需要使用三维张量来存储所有的解码结果。params.decode_batch_size
定义了每个进程生成的解码结果数量,pad_max
定义了解码结果的最大长度。
results
变量是一个空列表,用于存储最终的解码结果。
while True:
try:
features = next(iterator)
batch_size = features["source"].shape[0]
except:
features = {
"source": torch.ones([1, 1]).long(),
"source_mask": torch.ones([1, 1]).float()
}
batch_size = 0
t = time.time()
counter += 1
这是 _evaluate_model
函数的一部分,用于循环迭代数据集,对每个批次进行解码和填充,并计算处理时间。
使用 try
和 except
语句从数据集迭代器中获取下一个批次。如果没有下一个批次,则将 features
变量设置为包含一个张量和一个浮点数的字典,以表示一个空的批次。此外,将 batch_size
变量设置为 0。
然后,使用 time.time()
函数记录当前时间,并将其存储在 t
变量中。counter
变量用于记录迭代次数。
====
接下来,对批次进行解码和填充,并将结果存储在 seqs
变量中。具体来说,使用 beam_search
函数对批次进行解码,生成最佳解码序列。然后,使用 torch.squeeze
函数将解码后的结果张量的第二个维度压缩,以便将其转换为大小为 [batch_size, seq_len]
的二维张量。然后,使用 torch.nn.functional.pad
函数将解码后的结果填充为固定大小的张量,以便在分布式训练中进行合并。
然后,使用 size.zero_()
将 size
变量初始化为 0,使用 size[dist.get_rank()].copy_(torch.tensor(batch_size))
将当前进程的批次大小存储到 size
变量中,并使用 dist.all_reduce
函数进行同步。dist.all_reduce
函数将各个进程的 size
变量相加,并将结果存储在每个进程的 size
变量中。
接下来,使用 dist.all_gather
函数将所有进程的解码结果存储到 t_list
变量中。dist.all_gather
函数将 seqs
变量中的解码结果发送到所有进程,并将结果存储在 t_list
变量中。
如果批次大小为 0,则跳出循环。否则,将解码后的结果还原为原始数据,并将其添加到 results
列表中。
最后,计算处理时间,并打印一条消息,指示已完成的批次数和处理时间。
# Decode
seqs, _ = beam_search([model], features, params)
这是 _evaluate_model
函数的一部分,用于解码批次数据。
使用 beam_search
函数对批次数据进行解码,生成最佳解码序列。beam_search
函数接受三个参数:神经网络模型对象、输入数据和参数对象。在该函数中,将神经网络模型对象作为列表的形式传递给 beam_search
函数,这是因为 beam_search
函数可以同时对多个模型进行解码。
beam_search
函数使用束搜索算法生成最佳解码序列,其返回值是一个元组,包含两个张量:
seqs
:大小为[batch_size, seq_len]
的张量,表示每个批次中的最佳解码序列。scores
:大小为[batch_size]
的张量,表示每个批次中的最佳解码序列的得分。
在该函数中,只使用了 seqs
张量,而将 scores
张量忽略。
# Padding
seqs = torch.squeeze(seqs, dim=1)
pad_batch = params.decode_batch_size - seqs.shape[0]
pad_length = pad_max - seqs.shape[1]
seqs = torch.nn.functional.pad(seqs, (0, pad_length, 0, pad_batch))
这是 _evaluate_model
函数的一部分,用于对解码后的结果进行填充。
使用 torch.squeeze
函数将解码后的结果张量的第二个维度压缩,以便将其转换为大小为 [batch_size, seq_len]
的二维张量。然后,计算需要填充的批次数和填充长度,并使用 torch.nn.functional.pad
函数将解码后的结果填充为固定大小的张量。具体来说,torch.nn.functional.pad
函数使用 (0, pad_length, 0, pad_batch)
的元组参数对解码后的结果进行填充,其中 (0, pad_length)
表示在第二个维度上填充 pad_length
个 0,(0, pad_batch)
表示在第一个维度上填充 pad_batch
个 0。
# Synchronization
size.zero_()
size[dist.get_rank()].copy_(torch.tensor(batch_size))
dist.all_reduce(size)
dist.all_gather(t_list, seqs)
这是 _evaluate_model
函数的一部分,用于同步解码结果。
使用 size.zero_()
将 size
变量初始化为 0。然后,使用 size[dist.get_rank()].copy_(torch.tensor(batch_size))
将当前进程的批次大小存储到 size
变量中,并使用 dist.all_reduce
函数进行同步。dist.all_reduce
函数将各个进程的 size
变量相加,并将结果存储在每个进程的 size
变量中。
接下来,使用 dist.all_gather
函数将所有进程的解码结果存储到 t_list
变量中。dist.all_gather
函数将 seqs
变量中的解码结果发送到所有进程,并将结果存储在 t_list
变量中。这样,所有进程都可以访问解码结果,以便进行后续处理。
if size.sum() == 0:
break
if dist.get_rank() != 0:
continue
这是 _evaluate_model
函数的一部分,用于判断是否结束迭代循环。
使用 size.sum()
判断当前批次是否为空批次。如果是空批次,则跳出循环,结束解码过程。
如果当前进程的排名不是 0,则继续循环,等待进程 0 进行处理。这是因为进程 0 负责将所有进程的解码结果合并,因此需要等待进程 0 完成处理。
for i in range(params.decode_batch_size):
for j in range(dist.get_world_size()):
n = size[j]
seq = _convert_to_string(t_list[j][i], params)
if i >= n:
continue
# Restore BPE segmentation
seq = BPE.decode(seq)
results.append(seq.split())
t = time.time() - t
print("Finished batch: %d (%.3f sec)" % (counter, t))
这是 _evaluate_model
函数的一部分,用于将解码结果转换为字符串,并将其添加到 results
列表中。
使用嵌套的 for
循环遍历 t_list
变量中的所有解码结果。外层循环遍历解码结果的数量,内层循环遍历所有进程。在循环中,使用 size[j]
变量获取进程 j 中的批次大小,然后使用 _convert_to_string
函数将解码结果转换为字符串。如果当前解码结果的索引大于等于进程 j 中的批次大小,则跳过当前循环。
然后,使用 BPE.decode
函数还原解码结果的 BPE 分段,并使用 split
函数将其拆分为单个单词。最后,将拆分后的结果添加到 results
列表中。
打印一条消息,指示已完成的批次数和处理时间。
model.train()
if dist.get_rank() == 0:
restored_results = []
for idx in range(len(results)):
restored_results.append(results[sorted_key[idx]])
return bleu(restored_results, references)
return 0.0
这是 _evaluate_model
函数的一部分,用于计算 BLEU 分数。
使用 model.train()
将模型设置为训练模式。这是因为在分布式训练中,每个进程都需要独立进行训练。在进行 BLEU 分数计算之前,需要确保模型处于训练模式。
如果当前进程的排名是 0,则对解码结果进行排序,并使用 bleu
函数计算 BLEU 分数。具体来说,使用 sorted_key
对解码结果进行排序,然后将排序后的结果存储在 restored_results
变量中。最后,使用 bleu
函数计算排序后的结果和参考答案之间的 BLEU 分数,并将其返回。
如果当前进程的排名不是 0,则返回 0.0,表示该进程没有计算 BLEU 分数。在分布式训练中,只有进程 0 负责计算 BLEU 分数。
BLEU计算
from thumt.utils.bleu import bleu
def bleu(trans, refs, bp="closest", smooth=False, n=4, weights=None):
p_norm = [0 for _ in range(n)]
p_denorm = [0 for _ in range(n)]
for candidate, references in zip(trans, refs):
for i in range(n):
ccount, tcount = modified_precision(candidate, references, i + 1)
p_norm[i] += ccount
p_denorm[i] += tcount
bleu_n = [0 for _ in range(n)]
for i in range(n):
# add one smoothing
if smooth and i > 0:
p_norm[i] += 1
p_denorm[i] += 1
if p_norm[i] == 0 or p_denorm[i] == 0:
bleu_n[i] = -9999
else:
bleu_n[i] = math.log(float(p_norm[i]) / float(p_denorm[i]))
if weights:
if len(weights) != n:
raise ValueError("len(weights) != n: invalid weight number")
log_precision = sum([bleu_n[i] * weights[i] for i in range(n)])
else:
log_precision = sum(bleu_n) / float(n)
bp = brevity_penalty(trans, refs, bp)
score = bp * math.exp(log_precision)
return score
这是 bleu
函数的实现,用于计算 BLEU 分数。
该函数接受五个参数:
trans
:包含候选翻译的列表。refs
:包含参考翻译的列表。bp
:指定计算短文本惩罚的方法,包括 "closest"(默认值)和 "shortest"。smooth
:指定是否使用加一平滑(默认为 False)。n
:指定计算 BLEU 分数时要考虑的 n-gram 的最大长度(默认为 4)。weights
:包含用于加权平均 n-gram 分数的权重的列表。
在函数实现中,首先初始化 p_norm
和 p_denorm
列表,用于存储分子和分母的数量。然后,对于每个候选翻译和对应的参考翻译,使用 modified_precision
函数计算出每个 n-gram 的分子和分母的数量,并将它们添加到 p_norm
和 p_denorm
列表中。
接下来,初始化 bleu_n
列表,用于存储每个 n-gram 的对数精度。对于每个 n-gram,如果使用平滑并且 n > 0,则在分子和分母中添加一个计数。如果分子或分母中的任何一个值为 0,则将 bleu_n[i]
设置为 -9999,否则计算对数精度并将其添加到 bleu_n
列表中。
然后,如果提供了权重,则使用加权平均值计算对数精度,否则使用算术平均值。接下来,使用 brevity_penalty
函数计算短文本惩罚,并将其乘以对数精度的指数。最后,将得分返回。