Train

generation.Trainer

Defect Generator 모델을 학습시키기 위한 Trainer class 입니다.

학습 데이터 수 권장 사항

MODE	최소 이미지 수	비고
Pen	10장	-
Box	50장	10장으로 학습한 모델의 pen 모드로 생성한 이미지 (40장+) 활용 가능
Auto	50장	10장으로 학습한 모델의 pen 모드로 생성한 이미지 (40장+) 활용 가능

Usage

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

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

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

# 학습 루프: Defect Generator
error, message, _ = trainer.change_train_mode("defect")
assert error >= 0

for step in range(CONFIG["total_iterations"]):
    error, message, 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, message, result = trainer.validate()
        assert error >= 0
        print(f"Validation {step + 1}:", result)

# 학습: Defect Region Proposer
error, message, _ = trainer.train_auto_mode()
assert error >= 0

# 학습 루프: Mask Generator
error, message, _ = trainer.change_train_mode("mask")
assert error >= 0

for step in range(CONFIG["total_iterations"]):
    error, message, 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, message, result = trainer.validate()
        assert error >= 0
        print(f"Validation {step + 1}:", result)

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

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

build(config, data) classmethod

Trainer class의 instance를 생성합니다.

Parameters:

Name	Type	Description	Default
config	Dict	Trainer를 build 하기 위한 config가 담겨 있는 dictionary 입니다. { "network_type": str, # 사전 학습 모델로는 "v2_standard" 만을 지원합니다. "total_iterations": int, # total training iterations "n_classes": int, # number of data classes (including background class, n_classes >= 2) "roi": Optional[Dict], # roi 세팅. ROIHandlerAPI의 build config와 동일합니다. None이면 ROI 적용 안함. (default None) "crop_wh": Union[List[int], Tuple[int, int]], # crop size for cropper (default (256, 256)) # "resize_factor": Optional[float], # training image resize scale (default None) # XXX: not available for now "augmentation": Optional[List[Dict]], # 학습에 사용할 augmentation config들을 모은 list 입니다. (default None) "batch_size": int, # number of train batch size (default 4) "num_workers": int, # number of train data worker processes (default 4) # "inspection_size_wh": Optional[List[int]], # validation에서 사용할 inspection_size_wh 입니다. # XXX: not available for now }	required
data	Dict	Trainer가 사용 할 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": [ # 2가지 타입 지원: "bounding_box" + "bitmap" 타입 / "contours" 타입 { # "bounding_box" + "bitmap" 타입 "class_index": int, # class index (0: background) "bounding_box": List[int], # [left, top, width, height] "bitmap": str, # base64 encoded png image string (shape: (h_label, w_label)) }, { # "contours" 타입 "class_index": int, # class index (0: background) "contours": [ # cv2 형식의 contours. 0번째 contour가 outer, 나머지는 모두 inner. ndarray(int32, shape=(num_points_0, 1, 2)), ndarray(int32, shape=(num_points_1, 1, 2)), ndarray(int32, shape=(num_points_2, 1, 2)), ... # contour 개수 만큼 반복 ], }, ..., # 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
Trainer		build가 완료된 generation.Trainer class의 instance를 반환합니다.

to_device(device=None)

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

Parameters:

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

Returns: None: None

train_auto_mode()

오토 모드 사용을 위한 모듈들을 학습시킵니다.

change_train_mode(train_mode)

학습시킬 모듈을 변경합니다.

Parameters:

Name	Type	Description	Default
train_mode	str	["defect", "mask"] 중에서 선택가능합니다.	required

train_one_step()

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

Returns:

Name	Type	Description
Dict		학습 결과가 들어있는 dictionary를 반환합니다. train_mode "defect" 인 경우, { "step": int, # 현재까지 진행된 총 학습 step. "epoch": int, # 현재까지 진행된 총 학습 epoch. "epoch_step": int, # 현재 epoch에서 몇 번째 step인지를 나타냄. "loss/total": float, # 학습에 사용되는 모든 loss를 모두 더한 값. "loss/rec_loss/l1": float, "loss/rec_loss/style": float, "loss/rec_loss/perceptual": float, "loss/rec_loss/total_variance": float, "loss/gan_loss/generator": float, "loss/gan_loss/discriminator": float, "learning_rate/defect_generator/0": float, # optimizer의 learning rate "learning_rate/defect_discriminator/0": float, # optimizer의 learning rate "step_time(sec)": float, # 학습을 1 step 진행하는데 걸린 총 시간 "data_time(sec)": float, # 데이터 로드 및 전처리에 걸린 시간 "model_time(sec)": float, # 모델 학습에 걸린 시간 } train_mode "mask" 인 경우, { "step": int, # 현재까지 진행된 총 학습 step. "epoch": int, # 현재까지 진행된 총 학습 epoch. "epoch_step": int, # 현재 epoch에서 몇 번째 step인지를 나타냄. "loss/total": float, # 학습에 사용되는 모든 loss를 모두 더한 값. "loss/seg_loss": float, "learning_rate/mask_generator/0": float, # optimizer의 learning rate "step_time(sec)": float, # 학습을 1 step 진행하는데 걸린 총 시간 "data_time(sec)": float, # 데이터 로드 및 전처리에 걸린 시간 "model_time(sec)": float, # 모델 학습에 걸린 시간 }

validation_enabled()

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

Returns:

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

validate()

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

Returns:

Name	Type	Description
Dict		validation 결과가 들어있는 dictionary를 반환합니다. { "mae": float, "psnr": float, "F1": float, }

post_train_process()

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

Returns:

Name	Type	Description
Dict	Dict	각 parameter에 대한 설명은 DetRecInferencer.get_default_inference_option를 참고해주세요. { "mask.smooth_mask": bool, "mask.assist_intensity": int, }

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

완료된 학습 step을 저장하고 Trainer를 paused 상태로 전환합니다.

Generator와 모든 n_critic discriminator update가 끝날 때까지 기다립니다. Auto-mode fitting 도중 호출하면 fitting 전체가 끝날 때까지 기다립니다. 성공 응답을 받은 뒤 제품에서 프로세스를 종료할 수 있습니다. 두 generator와 활성 discriminator, 모든 optimizer/scheduler, step counter, defect/mask mode, network의 train/eval 상태, sampler 위치, 남은 crop buffer, RNG와 완료된 auto-mode region 설정을 함께 저장합니다. 일반 save_checkpoint() 형식은 유지하며 pause 이후 학습에는 resume()이 필요합니다.

학습 실패 시 누적 gradient를 비웁니다. Optimizer/scaler 갱신 단계 이후 실패하면 resume()으로 복원해야 다시 학습할 수 있습니다. Pause 성공 후 training loader를 닫습니다.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	pause checkpoint를 저장할 경로.	required
metadata	Optional[Dict]	외부 loop의 단계와 진행률 등 caller가 관리할 정보.	None
password	Optional[str]	checkpoint 암호화에 사용할 password.	None

Returns:

Name	Type	Description
None	None	성공 시 result는 None. Public API는 (error, message, result)를 반환하며 error >= 0인 경우에만 저장 성공입니다.

Note

SaigeToolkit >=2.8.3,<2.9.0이 필요합니다. 단일 training device를 지원합니다. Custom style augmentor는 지원하지 않습니다. Auto-mode fitting 내부의 부분 재개는 지원하지 않습니다. 저장 실패 시 이전 파일을 유지합니다. 성공으로 간주하거나 프로세스를 종료하지 마세요. 실패한 학습 step은 정상 step 완료 또는 이전 checkpoint 복원 전까지 pause할 수 없습니다.

resume(checkpoint_path, password=None)

Pause checkpoint를 복원하고 저장된 mode의 다음 step부터 학습합니다.

Module major/minor와 지원하는 training state format_version이 같아야 하며, patch 버전 차이는 양방향으로 허용합니다. 새 프로세스에서는 동일한 config/data로 build()하고 같은 device 종류를 to_device()로 선택한 뒤 호출합니다. 복원 후 change_train_mode()를 다시 호출하면 crop buffer가 초기화되므로 다음 train_one_step()을 호출합니다. 기존 save_checkpoint() 파일은 이 API로 resume할 수 없습니다.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	pause()로 저장한 checkpoint 경로.	required
password	Optional[str]	저장할 때 사용한 password.	None

Returns:

Name	Type	Description
Dict	Dict[str, Any]	저장한 metadata. 생략했다면 빈 dict입니다. Public API는 (error, message, metadata)를 반환하며 실패 시 result는 None입니다.

Note

데이터 내용·순서, transform, library version이 같아야 합니다. 기본 CPU augmentation의 global RNG를 복원하며 private RNG/stateful custom transform은 별도 상태 관리가 필요합니다. Fitting 이전 checkpoint로 되돌리면 이후에 생긴 auto-mode region 설정도 제거합니다.

Metadata의 stage와 completed_steps는 caller가 관리하는 예시 key입니다. 제품은 defect 학습 → auto-mode fitting → mask 학습의 외부 진행 단계를 저장해야 합니다. Trainer의 step counter는 두 학습 mode를 합친 값입니다.

Usage

def require_ok(response):
    error, message, result = response
    if error < 0:
        raise RuntimeError(f"[{error}] {message}")
    return result

# 프로세스 A: 학습 후 pause. 성공 응답을 받은 뒤 제품에서 종료합니다.
trainer = require_ok(Trainer.build(config=CONFIG, data=DATA))
require_ok(trainer.to_device(DEVICE))
require_ok(trainer.change_train_mode("defect"))
for _ in range(50):
    require_ok(trainer.train_one_step())
require_ok(trainer.pause(
    "training.pause",
    metadata={"stage": "defect", "completed_steps": 50},
    password=PASSWORD,
))

# 프로세스 B: 동일한 설정으로 build/to_device한 뒤 resume합니다.
trainer = require_ok(Trainer.build(config=CONFIG, data=DATA))
require_ok(trainer.to_device(DEVICE))
metadata = require_ok(trainer.resume("training.pause", password=PASSWORD))
for _ in range(metadata["completed_steps"], CONFIG["total_iterations"]):
    require_ok(trainer.train_one_step())

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

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

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로부터 Trainer의 상태를 불러오고, 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 입니다.