Skip to content

Train

core.Trainer

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

Usage

demo/train.py 참고

build(config, data) classmethod

Trainer class의 instance를 생성합니다.

Parameters:

Name Type Description Default
config dict

Trainer를 build 하기 위한 config가 담겨 있는 dictionary 입니다. 

{
    "network_type": str,                     # "v2a_fast" / "v2a_standard" / "v2a_accurate" 중 하나를 지원합니다.
    "total_iterations": int,                 # 총 학습 step 수 (default 10000)
    "image_resize_wh": list[int] | None,     # training image resize scale (default [512, 512])
    "roi": dict | None,                      # roi 세팅. ROIHandlerAPI의 build config와 동일합니다.
                                             # Trainer build시 설정한 ROI를 Validation, Inference시에도 적용됩니다.
                                             # None이면 ROI를 적용하지 않습니다. (default None)
    "augmentation": list[dict] | None,       # 학습에 사용할 augmentation config들을 모은 list 입니다. (default None)
    "num_workers": int,                      # number of train data worker processes (default 8)
    "batch_size": int,                       # train batch size (default 16)
}

required
data dict

Trainer가 사용 할 train/validation data가 담겨 있는 dictionary 입니다. 

{
    "train_images": {
        "{image_id}": {
            "path": str,  # image path
            "labels": {
                "is_ng": bool,  # ng or not
            },
        },
        ...,  # times number of images
    },
    "validation_images": {...},  # optional
}

required

Returns:

Name Type Description
Trainer Trainer

build가 완료된 core.Trainer class의 instance를 반환합니다.

to_device(device)

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

Parameters:

Name Type Description Default
device device | str | int

변경하고자 하는 device 입니다. int의 경우 해당 번호의 GPU, "cpu"(str)은 cpu를 사용합니다.

 required

train_one_step()

학습을 1스텝 수행합니다.

Returns:

Type Description
dict[str, Any]

dict[str, Any]: 훈련 1 스텝 후 리턴하는 값입니다. 

{
    "step": int,  # 현재까지 진행된 총 학습 step.
    "loss": float,  # 훈련 loss값.
    "step_time(sec)": float,  # 학습을 1 step 진행하는데 걸린 총 시간
    "data_time(sec)": float,  # 데이터를 로딩하는데 걸린 시간
    "model_time(sec)": float,  # 모델 학습에 걸린 시간
}

post_train_process()

train을 마무리하기 위한 작업을 합니다. 주로 threshold를 자동으로 계산해줍니다.

Returns:

Name Type Description
dict dict

각 parameter에 대한 설명은 InferenceHandler.get_default_inference_option를 참고해주세요. 

{
    "params.object_score_threshold": float,     # default: 0.0
    "params.object_area_threshold": int,        # default: 0
    "params.batch_size": int,                   # default: 1

    # post_train_process를 호출해서 바뀌는 값은 아래 4가지 입니다.
    "params.anomaly_score_threshold": float,    # default: 0.0
    "params.object_boundary_threshold": float,  # default: 0.0
    "params.heatmap_color_range_max": float,    # default: 0.0
    "params.heatmap_color_range_min": float,    # default: 0.0
}

validation_enabled()

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

Returns:

Name Type Description
bool bool

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

get_validation_schedule(total_iterations, density_level) classmethod

validation을 수행할 step을 반환합니다.

Parameters:

Name Type Description Default
total_iterations int

총 학습 스텝입니다.

 required
density_level int

validation을 수행할 step의 밀도입니다. [0, 1, 2] 중 1가지이며 작을수록 validation을 더 자주 수행합니다. 기본값은 1입니다.

 required

Returns:

Type Description
list[int]

list[int]: validation을 수행할 step이 담긴 list를 반환합니다.

validate()

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

Returns:

Name Type Description
dict dict

validation 결과가 들어있는 dictionary를 반환합니다. OK와 NG와 OK이미지가 모두 적어도 1개 이상 있을때 출력합니다. 아니면 validate을 할 수 없습니다. 

{
    "AUROC": float,  # Area under ROC
    "AUPR": float,  # Area under PR curve
    "Best threshold": float,  # Best threshold
    "Best F1-Score": float,  # F1 score with labels
}

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

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

Parameters:

Name Type Description Default
checkpoint_path str

checkpoint를 저장할 path 입니다.

 required
metadata dict | None

checkpoint에 저장할 추가 metadata. Defaults to None.

 None
password str | None

checkpoint 파일에서 중요한 정보를 암호화 하는데 사용되는 password 입니다. None이면 암호화하지 않습니다. Defaults to None.

 None

load_checkpoint(checkpoint_path, password=None)

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

Parameters:

Name Type Description Default
checkpoint_path str

load할 checkpoint가 저장 되어 있는 path 입니다.

 required
password str | None

checkpoint 파일에서 중요한 정보를 복호화 하는데 사용되는 password 입니다. None이면 복호화하지 않습니다. Defaults to None.

 None

Returns:

Name Type Description
dict dict

save_checkpoint에서 저장했던 metadata 입니다.