修复了bug：当样本中出现非常长的公式(对应的token数可能超过2048)，会导致给label进行embedding时index out of range
This commit is contained in:
三洋三洋
@@ -30,7 +30,8 @@ OCR_IMG_MAX_WIDTH = 768
|
||||
OCR_IMG_CHANNELS = 1 # 灰度图
|
||||
|
||||
# ocr模型训练数据集的最长token数
|
||||
MAX_TOKEN_SIZE = 2048 # 模型最长的embedding长度(默认512)
|
||||
MAX_TOKEN_SIZE = 1024 # 模型最长的embedding长度(默认512)
|
||||
# MAX_TOKEN_SIZE = 2048 # 模型最长的embedding长度(默认512)
|
||||
# MAX_TOKEN_SIZE = 600
|
||||
|
||||
# ocr模型训练时随机缩放的比例
|
||||
|
||||
@@ -9,13 +9,24 @@ from transformers import Trainer, TrainingArguments, Seq2SeqTrainer, Seq2SeqTrai
|
||||
|
||||
from .training_args import CONFIG
|
||||
from ..model.TexTeller import TexTeller
|
||||
from ..utils.functional import tokenize_fn, collate_fn, img_transform_fn
|
||||
from ..utils.functional import tokenize_fn, collate_fn, img_transform_fn, filter_fn
|
||||
from ..utils.metrics import bleu_metric
|
||||
from ...globals import MAX_TOKEN_SIZE, MIN_WIDTH, MIN_HEIGHT
|
||||
|
||||
|
||||
def train(model, tokenizer, train_dataset, eval_dataset, collate_fn_with_tokenizer):
|
||||
training_args = TrainingArguments(**CONFIG)
|
||||
debug_mode = True
|
||||
if debug_mode:
|
||||
training_args.auto_find_batch_size = False
|
||||
training_args.num_train_epochs = 2
|
||||
training_args.per_device_train_batch_size = 2
|
||||
training_args.per_device_eval_batch_size = 2 * training_args.per_device_train_batch_size
|
||||
training_args.jit_mode_eval = False
|
||||
training_args.torch_compile = False
|
||||
training_args.dataloader_num_workers = 1
|
||||
|
||||
|
||||
trainer = Trainer(
|
||||
model,
|
||||
training_args,
|
||||
@@ -34,7 +45,7 @@ def evaluate(model, tokenizer, eval_dataset, collate_fn):
|
||||
eval_config = CONFIG.copy()
|
||||
eval_config['predict_with_generate'] = True
|
||||
generate_config = GenerationConfig(
|
||||
max_new_tokens=MAX_TOKEN_SIZE,
|
||||
max_length=MAX_TOKEN_SIZE-200,
|
||||
num_beams=1,
|
||||
do_sample=False,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
@@ -67,23 +78,25 @@ if __name__ == '__main__':
|
||||
dataset = load_dataset(
|
||||
'/home/lhy/code/TexTeller/src/models/ocr_model/train/data/loader.py'
|
||||
)['train']
|
||||
dataset = dataset.filter(lambda x: x['image'].height > MIN_HEIGHT and x['image'].width > MIN_WIDTH)
|
||||
tokenizer = TexTeller.get_tokenizer('/home/lhy/code/TexTeller/src/models/tokenizer/roberta-tokenizer-7Mformulas')
|
||||
filter_fn_with_tokenizer = partial(filter_fn, tokenizer=tokenizer)
|
||||
|
||||
# dataset = dataset.filter(lambda x: x['image'].height > MIN_HEIGHT and x['image'].width > MIN_WIDTH)
|
||||
dataset = dataset.filter(filter_fn_with_tokenizer, num_proc=16)
|
||||
dataset = dataset.shuffle(seed=42)
|
||||
dataset = dataset.flatten_indices()
|
||||
|
||||
tokenizer = TexTeller.get_tokenizer('/home/lhy/code/TeXify/src/models/tokenizer/roberta-tokenizer-550Kformulas')
|
||||
|
||||
map_fn = partial(tokenize_fn, tokenizer=tokenizer)
|
||||
tokenized_dataset = dataset.map(map_fn, batched=True, remove_columns=dataset.column_names, num_proc=8, load_from_cache_file=True)
|
||||
tokenized_dataset = tokenized_dataset.with_transform(img_transform_fn)
|
||||
|
||||
split_dataset = tokenized_dataset.train_test_split(test_size=0.05, seed=42)
|
||||
split_dataset = tokenized_dataset.train_test_split(test_size=0.005, seed=42)
|
||||
train_dataset, eval_dataset = split_dataset['train'], split_dataset['test']
|
||||
collate_fn_with_tokenizer = partial(collate_fn, tokenizer=tokenizer)
|
||||
model = TexTeller()
|
||||
# model = TexTeller.from_pretrained('/home/lhy/code/TeXify/src/models/ocr_model/train/train_result/train_with_random_resize/checkpoint-80000')
|
||||
|
||||
enable_train = False
|
||||
enable_train = True
|
||||
enable_evaluate = True
|
||||
if enable_train:
|
||||
train(model, tokenizer, train_dataset, eval_dataset, collate_fn_with_tokenizer)
|
||||
|
||||
@@ -4,7 +4,7 @@ CONFIG = {
|
||||
# "data_seed": 42, # data sampler的采样也固定
|
||||
# "full_determinism": True, # 使整个训练完全固定（这个设置会有害于模型训练，只用于debug）
|
||||
|
||||
"output_dir": "train_result/train_with_random_resize", # 输出目录
|
||||
"output_dir": "train_result/TexTellerv2", # 输出目录
|
||||
"overwrite_output_dir": False, # 如果输出目录存在，不删除原先的内容
|
||||
"report_to": ["tensorboard"], # 输出日志到TensorBoard，
|
||||
#+通过在命令行：tensorboard --logdir ./logs 来查看日志
|
||||
@@ -12,7 +12,7 @@ CONFIG = {
|
||||
"logging_dir": None, # TensorBoard日志文件的存储目录(使用默认值)
|
||||
"log_level": "warning", # 其他可选:‘debug’, ‘info’, ‘warning’, ‘error’ and ‘critical’（由低级别到高级别）
|
||||
"logging_strategy": "steps", # 每隔一定步数记录一次日志
|
||||
"logging_steps": 500, # 记录日志的步数间隔，可以是int也可以是(0~1)的float，当是float时表示总的训练步数的ratio(比方说可以设置成1.0 / 2000)
|
||||
"logging_steps": 4000, # 记录日志的步数间隔，可以是int也可以是(0~1)的float，当是float时表示总的训练步数的ratio(比方说可以设置成1.0 / 2000)
|
||||
#+通常与eval_steps一致
|
||||
"logging_nan_inf_filter": False, # 对loss=nan或inf进行记录
|
||||
|
||||
@@ -25,6 +25,7 @@ CONFIG = {
|
||||
"per_device_train_batch_size": 64, # 每个GPU的batch size
|
||||
"per_device_eval_batch_size": 16, # 每个GPU的evaluation batch size
|
||||
"auto_find_batch_size": True, # 自动搜索合适的batch size（指数decay）
|
||||
# "auto_find_batch_size": False, # 自动搜索合适的batch size（指数decay）
|
||||
|
||||
"optim": "adamw_torch", # 还提供了很多AdamW的变体（相较于经典的AdamW更加高效）
|
||||
#+当设置了optim后，就不需要在Trainer中传入optimizer
|
||||
@@ -52,12 +53,12 @@ CONFIG = {
|
||||
"dataloader_drop_last": True, # 丢掉最后一个minibatch，保证训练的梯度稳定
|
||||
|
||||
"evaluation_strategy": "steps", # 评估策略，可以是"steps"或"epoch"
|
||||
"eval_steps": 500, # if evaluation_strategy="step"
|
||||
"eval_steps": 4000, # if evaluation_strategy="step"
|
||||
#+默认情况下与logging_steps一样，可以是int也可以是(0~1)的float，当是float时表示总的训练步数的ratio(比方说可以设置成1.0 / 2000)
|
||||
|
||||
"save_strategy": "steps", # 保存checkpoint的策略
|
||||
"save_steps": 500, # checkpoint保存的步数间隔，可以是int也可以是(0~1)的float，当是float时表示总的训练步数的ratio(比方说可以设置成1.0 / 2000)
|
||||
"save_total_limit": 5, # 保存的模型的最大数量。如果超过这个数量，最旧的模型将被删除
|
||||
"save_steps": 4000, # checkpoint保存的步数间隔，可以是int也可以是(0~1)的float，当是float时表示总的训练步数的ratio(比方说可以设置成1.0 / 2000)
|
||||
"save_total_limit": 10, # 保存的模型的最大数量。如果超过这个数量，最旧的模型将被删除
|
||||
|
||||
"load_best_model_at_end": True, # 训练结束时是否加载最佳模型
|
||||
#+当设置True时，会保存训练时评估结果最好的checkpoint
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from functools import partial
|
||||
from datasets import load_dataset
|
||||
|
||||
from transformers import DataCollatorForLanguageModeling
|
||||
from typing import List, Dict, Any
|
||||
from ..model.TexTeller import TexTeller
|
||||
from .transforms import train_transform
|
||||
from ..model.TexTeller import TexTeller
|
||||
from ...globals import MIN_HEIGHT, MIN_WIDTH, MAX_TOKEN_SIZE
|
||||
|
||||
|
||||
def left_move(x: torch.Tensor, pad_val):
|
||||
@@ -50,6 +50,13 @@ def img_transform_fn(samples: Dict[str, List[Any]]) -> Dict[str, List[Any]]:
|
||||
return samples
|
||||
|
||||
|
||||
def filter_fn(sample, tokenizer=None) -> bool:
|
||||
return (
|
||||
sample['image'].height > MIN_HEIGHT and sample['image'].width > MIN_WIDTH
|
||||
and len(tokenizer(sample['latex_formula'])['input_ids']) < MAX_TOKEN_SIZE - 10
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset = load_dataset(
|
||||
'/home/lhy/code/TeXify/src/models/ocr_model/train/dataset/latex-formulas/latex-formulas.py',
|
||||
|
||||
Reference in New Issue
Block a user