Skip to content

Train

generation.Trainer

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

학습 데이터 수 권장 사항
MODE 최소 이미지 수 비고
Pen 10장 -
Box 50장 10장으로 학습한 모델의 pen 모드로 생성한 이미지 (40장+) 활용 가능
Auto 50장 10장으로 학습한 모델의 pen 모드로 생성한 이미지 (40장+) 활용 가능
Usage
# Trainer 빌드
error, trainer = Trainer.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

# 학습 루프: Defect Generator
error, _ = trainer.change_train_mode("defect")
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)

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

# 학습 루프: Mask Generator
error, _ = trainer.change_train_mode("mask")
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

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,
    "ssim": 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,
}

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 입니다.