KeyInformationExtraction

ocr.all.KieTrainer

Key Information Extraction 모델을 학습시키기 위한 Trainer class 입니다.

Usage

# Trainer 빌드
error, trainer = KieTrainer.build(CONFIG, DATA)
assert error >= 0

# 0번 GPU로 이동
error, result = trainer.to_device(DEVICE)
assert error >= 0

# validation 가능한지 확인
error, validation_enabled = trainer.validation_enabled()
assert error >= 0

# 학습 루프
for step in range(CONFIG["total_iterations"]):
    error, result = trainer.train_one_step()
    assert error >= 0

    if (step + 1) % PRINT_INTERVAL == 0:
        print(f"Train {step + 1}:", result)

    if validation_enabled and (step + 1) % VALIDATION_INTERVAL == 0:
        error, result = trainer.validate()
        assert error >= 0
        print(f"Validation {step + 1}:", result)

# 학습 후처리 (default inference options)
error, post_train_process_result = trainer.post_train_process()
assert error >= 0
print("post_train_process :", post_train_process_result)

# 체크포인트 저장
error, _ = trainer.save_checkpoint(
    checkpoint_path=CHECKPOINT_PATH,
    password=PASSWORD,
    metadata={"description": "demo"},
)

build(config, data) classmethod

KieTrainer class의 instance를 생성합니다.

Parameters:

Name	Type	Description	Default
config	Dict	KieTrainer를 build 하기 위한 config가 담겨 있는 dictionary 입니다. { "network_type": str, # "v2_cj" 중 하나를 지원합니다. "total_iterations": int, # total training iterations "roi": Optional[Dict], # roi 세팅. ROIHandlerAPI의 build config와 동일합니다. None이면 ROI 적용 안함. (default None) "target_text_size": Optional[int], # 리사이즈 후 target text size (16-64, default None). 하단 설명 참고. "augmentation": Optional[List[Dict]], # 학습에 사용할 augmentation config들을 모은 list 입니다. (default None) "batch_size": int, # number of train batch size (default 8) "num_workers": int, # number of train data worker processes (default 4) }	required
data	Dict	KieTrainer가 사용 할 train/validation data가 담겨 있는 dictionary 입니다. { "train_images": { # 1: basic structure "{image_id}": { "path": Union[str, List[str]], # 이미지 경로 (multipage인 경우 경로 리스트) "width": int, # image width "height": int, # image height "labels": [ { "polygon": [ [point_1_x, point_1_y], # TODO: 폴리곤 정렬 순서에 대해 제품팀과 논의 [point_2_x, point_2_y], ..., [point_n_x, point_n_y] ], # float "text": text in box # str "class_index": int, # class index. XXX: Not implemented yet }, ...# times number of labels in the image ], }, # 2: label path structure "{image_id}": { "path": Union[str, List[str]], # 이미지 경로 (multipage인 경우 경로 리스트) "width": int, # image width "height": int, # image height "labels_count": int, # len(labels) "labels_path": str, # pickle path to "labels" }, ..., # times number of images }, "validation_images": {...}, }	required

Returns:

Name	Type	Description
KieTrainer		build가 완료된 ocr.KieTrainer class의 instance를 반환합니다.

to_device(device=None)

KieTrainer의 device를 변경합니다. (cpu/gpu)

Parameters:

Name	Type	Description	Default
device	Union[device, str, int]	변경하고자 하는 device 입니다. int의 경우 해당 번호의 GPU를, "cpu" 문자열의 경우 cpu를 사용합니다.	None

Returns: None: None

train_one_step()

학습을 1스텝 수행하고, 결과를 반환합니다.

Returns:

Name	Type	Description
Dict		학습 결과가 들어있는 dictionary를 반환합니다. { "step": int, # 현재까지 진행된 총 학습 step. "epoch": int, # 현재까지 진행된 총 학습 epoch. "epoch_step": int, # 현재 epoch에서 몇 번째 step인지를 나타냄. "loss/total": float, # 학습에 사용되는 모든 loss를 모두 더한 값. "loss/cross_entropy": float, "learning_rate/0": float, # optimizer의 learning rate "step_time(sec)": float, # 학습을 1 step 진행하는데 걸린 총 시간 "data_time(sec)": float, # 데이터를 로딩하는데 걸린 시간 "model_time(sec)": float, # 모델 학습에 걸린 시간 }

validation_enabled()

KieTrainer가 validation이 가능한지 bool로 반환합니다.

Returns:

Name	Type	Description
bool		KieTrainer가 validation이 가능하면 True를, 불가능하다면 False를 반환합니다.

validate()

Validation을 수행하고, 결과를 리턴합니다.

Returns:

Name	Type	Description
Dict		validation 결과가 들어있는 dictionary를 반환합니다. # NOTE: "sim_value" is key metric for best model. { "acc_total": float, # total accuracy "acc_key": float, # key class accuracy "acc_value": float, # value class accuracy "sim_total": float, # total similarity "sim_key": float, # key class similarity "sim_value": float, # value class similarity }

post_train_process()

train 후처리를 수행하고, 해당 모델의 default inference options을 리턴합니다.

Returns:

Name	Type	Description
Dict	Dict	각 parameter에 대한 설명은 DetRecKieInferencer.get_default_inference_option를 참고해주세요. {}

pause(checkpoint_path, metadata=None, password=None)

Continuous Training: 완료된 step의 상태를 저장하고 학습을 일시 정지합니다.

SaigeToolkit >=2.8.3,<2.9.0이 필요합니다. 저장한 파일은 같은 Trainer 또는 새 프로세스의 KieTrainer.resume()으로 복원합니다.

저장되는 상태

• 학습: network weight, optimizer/scheduler, step·epoch·epoch_step.
• 데이터: sampler의 실제 소비 위치와 worker prefetch를 제외한 소비 이력.
• 실행: Python·NumPy·Torch RNG와 network 및 child module의 train/eval 상태.
• Validation: warm-up 상태와 persistent worker iterator 생성 여부.

중단 동작

• 진행 중인 train_one_step()이 있으면 해당 step이 끝날 때까지 기다립니다.
• 저장 성공 후 training worker를 종료하며, resume() 전까지 추가 학습을 차단합니다.
• 프로세스 종료는 저장 성공을 확인한 caller가 처리합니다.

지원 조건

• 단일 training device의 MemorySafeDataLoader와 OCRFull2PatchDataLoader를 지원합니다.
• 지원 batch sampler는 naive, class_balance, balance입니다.
• 미지원 sampler는 오류 code -650336, multi-GPU는 -650337을 반환합니다.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	학습 재개용 checkpoint를 저장할 경로입니다.	required
metadata	Optional[Dict]	caller가 함께 저장할 추가 정보입니다. 생략하면 빈 dictionary를 저장합니다. Defaults to None.	None
password	Optional[str]	checkpoint 암호화에 사용할 password입니다. None이면 암호화하지 않습니다. Defaults to None.	None

반환값

성공 시 (0, message, None)을 반환합니다. 실패 시 (음수 error code, message, None)을 반환합니다.

Usage

학습 중인 Trainer의 상태를 저장하는 예제입니다.

code, message, _ = trainer.pause(
    "training.pause.pt", metadata={"phase": "training"}
)
assert code == 0, message
# 성공을 확인한 뒤 프로세스를 종료할 수 있습니다.

새 프로세스에서 이어서 학습하는 순서는 resume() 예제를 참고합니다.

resume(checkpoint_path, password=None)

Pause checkpoint를 복원하고 다음 batch부터 학습을 이어갑니다.

복원 순서

1. 저장할 때와 같은 config와 data로 KieTrainer.build()를 호출합니다.
2. to_device()로 저장할 때와 같은 device 종류를 선택합니다.
3. resume() 성공을 확인한 뒤 train_one_step()을 호출합니다.

호환성과 지원 범위

• Module major/minor가 같아야 합니다. Patch 차이는 양방향으로 허용합니다.
• Config, 데이터 내용·순서, library version과 device 종류를 유지해야 합니다. Synthetic data에 사용하는 font·background 등 외부 자료도 같아야 합니다.
• Worker augmentation의 global Python·NumPy·Torch CPU RNG를 복원합니다. Custom dataset/transform의 private RNG, CUDA transform, 별도 mutable state와 stochastic validation worker의 내부 진행 상태는 지원하지 않습니다.
• 일반 save_checkpoint() 파일에는 재개용 loader/RNG 상태가 없습니다. 학습 재개에는 pause()로 저장한 파일을 사용해야 합니다.
• Optimizer 갱신 중 학습에 실패했다면, 이전에 성공한 pause checkpoint를 resume()으로 복원한 뒤 학습을 다시 시작합니다.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	pause()로 저장한 checkpoint 경로입니다.	required
password	Optional[str]	저장할 때 사용한 password입니다. 암호화하지 않았다면 None입니다. Defaults to None.	None

반환값

성공 시 (0, message, metadata)를 반환합니다. Metadata를 생략하고 저장했다면 빈 dictionary입니다.

실패 시 (음수 error code, message, None)을 반환합니다.

Usage

같은 config와 data를 준비한 새 프로세스에서 학습을 재개하는 예제입니다. Device는 pause할 때 사용한 device 종류에 맞춰 선택합니다.

from ocr.all import KieTrainer

code, message, trainer = KieTrainer.build(config, data)
assert code == 0, message
code, message, _ = trainer.to_device("cuda:0")
assert code == 0, message
code, message, metadata = trainer.resume("training.pause.pt")
assert code == 0, message
code, message, result = trainer.train_one_step()
assert code == 0, message

같은 Trainer에서 pause()한 뒤 다시 resume()할 수도 있습니다. 암호화한 checkpoint는 pause()와 같은 password를 전달해야 합니다.

save_checkpoint(checkpoint_path, metadata=None, password=None)

현재 KieTrainer의 상태를 암호화하여 저장합니다.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	checkpoint를 저장할 path 입니다.	required
metadata	Optional[Dict]	checkpoint에 저장할 추가 metadata. Defaults to None.	None
password	Optional[str]	checkpoint 파일에서 중요한 정보를 암호화 하는데 사용되는 password 입니다. None이면 암호화하지 않습니다. Defaults to None.	None

Returns: None: None

load_checkpoint(checkpoint_path, password=None, **options)

저장한 checkpoint로부터 KieTrainer의 상태를 불러오고, metadata를 반환합니다.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	load할 checkpoint가 저장 되어 있는 path 입니다.	required
password	Optional[str]	checkpoint 파일에서 중요한 정보를 복호화 하는데 사용되는 password 입니다. None이면 복호화하지 않습니다. Defaults to None.	None

Returns:

Name	Type	Description
Dict	dict	save_checkpoint에서 저장했던 metadata 입니다.