콘텐츠로 이동

Anomaly Clustering Train

vims.AnomalyClusteringTrainer

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

Usage

demo/train_anomaly_cluster.py 참고

build(config, data) classmethod

AnomalyClusteringTrainer class의 instance를 생성합니다.

Parameters:

Name Type Description Default
config Dict

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

{
    "checkpoint_path": str, 기존 vims 모델의 checkpoint path
    "password": str,  vims 모델의 checkpoint path
    "clustering_method": Literal['ssim'], clustering 방법입니다. default는 'ssim' 입니다.
    "device": Union[int, str], GPU 번호입니다. default는 '0' 입니다.
    "pretrained_model_weight_dir": str, pretrained model weight가 저장된 directory 입니다. default는 'vims/pretrained_weight/clustering' 입니다.
}

required
data List[dict]

AnomalyClusteringTrainer가 사용 할 train data가 담겨 있는 dictionary list 입니다.

[
    {
        "img_path": str,
    },
]

required

train()

AnomalyClusteringTrainer를 통해 학습을 수행합니다.

Returns:

Type Description
List[dict]

List[dict]: 주어진 이미지에 대해서, 학습 결과가 담긴 dictionary list 입니다.

[
    {
        "img_path": str,
        "cluster_index": int,
    },
]

to_device(device)

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

Parameters:

Name Type Description Default
device Union[device, str, int]

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

required

Returns:

Name Type Description
None None

None

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:

Name Type Description
None None

None