Skip to content

Train

enhancement.UnpairedTrainer

Enhancement 모델을 학습시키기 위한 UnpairedTrainer class 입니다.

Usage

demo/train.py 참고

build(config, data) classmethod

UnpairedTrainer class의 instance를 생성합니다.

Parameters:

Name Type Description Default
config Dict

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

{
    "network_type": str,  # network type. "v2_fast" / "v2_standard" 중 하나를 지원합니다.
    "total_iterations": int,  # 총 학습 iteration 수
    "crop_wh": int,  # 학습에 사용할 이미지 크기 (default (256, 256))
    "num_workers": int,  # 학습에 사용할 dataloader의 num_workers (default 6)
    "inspection_size_wh": Optional[List[int]],  # post_train_process에서 사용할 inspection_size_wh 입니다.
                                                # 자세한 정보는 InferenceHandler의 get_default_inference_option 함수 설명에 나와있습니다.
}

required
data Dict

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

{
    "low_images": {
        "{image_id}": {
            "path": str,  # 저품질 이미지 경로
        },
        ...,  # times number of images
    },
    "high_images": {
        "{image_id}": {
            "path": str,  # 고품질 이미지 경로
        },
        ...,  # times number of images
    },
}

required

Returns:

Name Type Description
UnpairedTrainer UnpairedTrainer

build가 완료된 enhancement.UnpairedTrainer class의 instance를 반환합니다.

to_device(device)

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

Parameters:

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

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

required

Returns: None: None

train_one_step()

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

Returns:

Name Type Description
Dict Dict

학습 결과가 들어있는 dictionary를 반환합니다.

{
    "step": int,  # 현재까지 진행된 총 학습 step.
    "epoch": int,  # 현재까지 진행된 총 학습 epoch.
    "epoch_step": int,  # 현재 epoch에서 몇 번째 step인지를 나타냄.
    "loss/G": float,  # generator의 총 loss.
    "loss/D_A": float,  # discriminator A의 loss.
    "loss/D_B": float,  # discriminator B의 loss.
    "loss/G_A": float,  # generator A의 loss.
    "loss/G_B": float,  # generator B의 loss.
    "loss/cycle_A": float,  # cycle loss A.
    "loss/cycle_B": float,  # cycle loss B.
    "loss/idt_A": float,  # identity loss A.
    "loss/idt_B": float,  # identity loss B.
    "learning_rate/0": float,  # optimizer 1의 learning rate.
    "learning_rate/1": float,  # optimizer 2의 learning rate.
    "step_time(sec)": float,  # 학습을 1 step 진행하는데 걸린 총 시간.
    "data_time(sec)": float,  # 데이터를 로딩하는데 걸린 시간.
    "model_time(sec)": float,  # 모델 학습에 걸린 시간.
}

validation_enabled()

Trainer가 validation이 가능한지 bool로 반환합니다. 현재 IEN은 validation을 지원하지 않습니다.

Returns:

Name Type Description
bool bool

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

post_train_process(save_root)

PairedTrainer의 학습에 사용할 데이터를 생성하고, PairedTrainer의 build에 필요한 data config를 리턴합니다.

Parameters:

Name Type Description Default
save_root str

학습에 사용할 데이터를 저장할 root path 입니다.

required

Returns:

Name Type Description
Dict Dict

PairedTrainer의 build에 필요한 data config가 담겨 있는 dictionary 입니다.

{
    "train_images": {
        "{image_id}": {
            "path_low": str,  # 저품질 이미지 경로
            "path_high": str,  # 고품질 이미지 경로
        },
        ...,  # times number of images
    },
}

Usage

post_train_process로 PairedTrainer의 학습에 사용할 데이터를 생성하여 build 하는 예시입니다. 상세 내용은 demo/train.py를 참고해주세요.

# PairedTrainer 학습에 사용할 데이터 생성
error, message, paired_data = unpaired_trainer.post_train_process(UNPAIRED_TRAIN_RESULTS_PATH)

# PairedTrainer 빌드
error, message, paired_trainer = PairedTrainer.build(PAIRED_CONFIG, paired_data)

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

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

Model, 모든 optimizer/scheduler, step/epoch/epoch_step, training loader 위치, RNG, validation warm-up과 persistent worker의 iterator 생성 여부를 저장합니다. CycleGAN의 과거 생성 이미지를 담는 두 image pool도 저장합니다. 성공한 뒤에는 resume() 전까지 추가 학습할 수 없습니다. 저장 성공을 확인한 뒤 프로세스를 종료하는 것은 제품에서 처리합니다.

SaigeToolkit >=2.8.3,<2.9.0과 use_infinite_sampler=True인 training loader가 필요합니다. Public build()는 이 설정을 사용합니다. Finite loader는 기존 학습을 지원하며, pause()/resume()은 오류 code -416336을 반환합니다. 일반 save_checkpoint() 형식은 유지되며, 학습 재개에는 이 API로 저장한 파일이 필요합니다.

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

Parameters:

Name Type Description Default
checkpoint_path str

재개용 checkpoint를 저장할 path입니다.

required
metadata Optional[Dict[str, Any]]

함께 저장할 caller의 추가 정보입니다. 생략하면 빈 dictionary를 저장합니다. Defaults to None.

None
password Optional[str]

checkpoint 암호화에 사용할 password입니다. None이면 암호화하지 않습니다. Defaults to None.

None

Returns:

Name Type Description
None None

성공 시 (error, message, None)을 반환합니다. 실패하면 error가 음수입니다.

Usage

error, message, _ = trainer.pause(
    "training.pause.pt", metadata={"completed_steps": completed_steps}, password=password
)
assert error >= 0, message
completed_steps는 caller가 정한 예시 key입니다. 내부 학습 counter는 별도로 저장합니다.

resume(checkpoint_path, password=None)

Pause checkpoint를 복원하여 다음 step부터 학습하고 저장된 metadata를 반환합니다.

Module major/minor와 지원하는 training state format_version이 같아야 하며, patch 버전 차이는 양방향으로 허용합니다. 같은 config, 데이터 내용과 순서, transform, library version을 사용해야 합니다. 새 프로세스에서는 같은 Trainer를 build하고 저장할 때와 같은 device 종류를 to_device()로 선택한 뒤 호출합니다. 같은 Trainer에서 재개할 수도 있습니다. 일반 checkpoint는 기존 load_checkpoint()로 읽습니다.

Training augmentation은 global Python/NumPy/Torch CPU RNG를 복원합니다. Private RNG, CUDA transform, 상태가 변하는 dataset은 별도 state 관리가 필요합니다. Persistent validation worker 내부의 RNG는 저장하지 않으므로 해당 validation transform은 결정적이어야 합니다. Iterator 생성 여부가 없는 이전 pause 파일은 읽을 수 있지만 validation worker lifecycle에 따른 RNG 일치는 보장하지 않습니다.

Parameters:

Name Type Description Default
checkpoint_path str

pause()로 저장한 checkpoint의 path입니다.

required
password Optional[str]

저장할 때 사용한 password입니다. 암호화하지 않았다면 None입니다. Defaults to None.

None

Returns:

Type Description
Dict[str, Any]

Dict[str, Any]: 성공 시 (error, message, metadata)를 반환합니다. Metadata를 생략하고 저장했다면 빈 dictionary입니다. 실패하면 음수 error와 None payload를 반환합니다.

Usage

새 프로세스에서 학습을 재개하는 예시입니다.

error, message, trainer = UnpairedTrainer.build(config, data)
assert error >= 0, message
error, message, _ = trainer.to_device(device)
assert error >= 0, message
error, message, metadata = trainer.resume("training.pause.pt", password=password)
assert error >= 0, message
# pause() 예제에서 caller가 저장한 진행률을 사용합니다.
for step in range(metadata["completed_steps"], total_steps):
    error, message, result = trainer.train_one_step()
    assert error >= 0, message

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)

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

enhancement.PairedTrainer

Enhancement 모델을 학습시키기 위한 PairedTrainer class 입니다.

Usage

demo/train.py 참고

build(config, data) classmethod

PairedTrainer class의 instance를 생성합니다.

Parameters:

Name Type Description Default
config Dict

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

{
    "network_type": str,  # network type. "v2_fast" / "v2_standard" 중 하나를 지원합니다.
    "total_iterations": int,  # 총 학습 iteration 수
    "batch_size": int,  # number of train batch size (default 16)
    "crop_wh": int,  # 학습에 사용할 이미지 크기 (default (256, 256))
    "augmentation": Optional[List[Dict]],  # 학습에 사용할 augmentation config들을 모은 list 입니다. (default None)
    "num_workers": int,  # 학습에 사용할 dataloader의 num_workers (default 6)
    "inspection_size_wh": Optional[List[int]],  # validation에서 사용할 inspection_size_wh 입니다.
                                                # 자세한 정보는 InferenceHandler의 get_default_inference_option 함수 설명에 나와있습니다.
}

required
data Dict

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

{
    "train_images": {
        "{image_id}": {
            "path_low": str,  # 저품질 이미지 경로
            "path_high": str,  # 고품질 이미지 경로
        },
        ...,  # times number of images
    }
}

required

Returns:

Name Type Description
PairedTrainer PairedTrainer

build가 완료된 enhancement.PairedTrainer class의 instance를 반환합니다.

to_device(device)

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

Parameters:

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

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

required

Returns: None: None

train_one_step()

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

Returns:

Name Type Description
Dict Dict

학습 결과가 들어있는 dictionary를 반환합니다.

{
    "step": int,  # 현재까지 진행된 총 학습 step.
    "epoch": int,  # 현재까지 진행된 총 학습 epoch.
    "epoch_step": int,  # 현재 epoch에서 몇 번째 step인지를 나타냄.
    "loss/total": float,  # 학습에 사용되는 모든 loss를 모두 더한 값.
    "loss/L1Loss": float,  # l1 loss.
    "loss/VGGLoss": float,  # vgg loss.
    "learning_rate/0": float,  # optimizer의 learning rate.
    "step_time(sec)": float,  # 학습을 1 step 진행하는데 걸린 총 시간.
    "data_time(sec)": float,  # 데이터를 로딩하는데 걸린 시간.
    "model_time(sec)": float,  # 모델 학습에 걸린 시간.
}

validation_enabled()

Trainer가 validation이 가능한지 bool로 반환합니다. 현재 IEN은 validation을 지원하지 않습니다.

Returns:

Name Type Description
bool bool

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

post_train_process()

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

Returns:

Name Type Description
Dict Dict

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

{
    "params.batch_size": int,
    "params.inspection_size_wh": Optional[List[int]],
    "params.oversized_image_handling": str,
}

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

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

Model, optimizer/scheduler, step/epoch/epoch_step, training loader 위치, RNG, validation warm-up과 persistent worker의 iterator 생성 여부를 저장합니다. 성공한 뒤에는 resume() 전까지 추가 학습할 수 없습니다. 저장 성공을 확인한 뒤 프로세스를 종료하는 것은 제품에서 처리합니다.

SaigeToolkit >=2.8.3,<2.9.0과 use_infinite_sampler=True인 training loader가 필요합니다. Public build()는 이 설정을 사용합니다. Finite loader는 기존 학습을 지원하며, pause()/resume()은 오류 code -416336을 반환합니다. 일반 save_checkpoint() 형식은 유지되며, 학습 재개에는 이 API로 저장한 파일이 필요합니다.

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

Parameters:

Name Type Description Default
checkpoint_path str

재개용 checkpoint를 저장할 path입니다.

required
metadata Optional[Dict[str, Any]]

함께 저장할 caller의 추가 정보입니다. 생략하면 빈 dictionary를 저장합니다. Defaults to None.

None
password Optional[str]

checkpoint 암호화에 사용할 password입니다. None이면 암호화하지 않습니다. Defaults to None.

None

Returns:

Name Type Description
None None

성공 시 (error, message, None)을 반환합니다. 실패하면 error가 음수입니다.

Usage

error, message, _ = trainer.pause(
    "training.pause.pt", metadata={"completed_steps": completed_steps}, password=password
)
assert error >= 0, message
completed_steps는 caller가 정한 예시 key입니다. 내부 학습 counter는 별도로 저장합니다.

resume(checkpoint_path, password=None)

Pause checkpoint를 복원하여 다음 step부터 학습하고 저장된 metadata를 반환합니다.

Module major/minor와 지원하는 training state format_version이 같아야 하며, patch 버전 차이는 양방향으로 허용합니다. 같은 config, 데이터 내용과 순서, transform, library version을 사용해야 합니다. 새 프로세스에서는 같은 Trainer를 build하고 저장할 때와 같은 device 종류를 to_device()로 선택한 뒤 호출합니다. 같은 Trainer에서 재개할 수도 있습니다. 일반 checkpoint는 기존 load_checkpoint()로 읽습니다.

Training augmentation은 global Python/NumPy/Torch CPU RNG를 복원합니다. Private RNG, CUDA transform, 상태가 변하는 dataset은 별도 state 관리가 필요합니다. Persistent validation worker 내부의 RNG는 저장하지 않으므로 해당 validation transform은 결정적이어야 합니다. Iterator 생성 여부가 없는 이전 pause 파일은 읽을 수 있지만 validation worker lifecycle에 따른 RNG 일치는 보장하지 않습니다.

Parameters:

Name Type Description Default
checkpoint_path str

pause()로 저장한 checkpoint의 path입니다.

required
password Optional[str]

저장할 때 사용한 password입니다. 암호화하지 않았다면 None입니다. Defaults to None.

None

Returns:

Type Description
Dict[str, Any]

Dict[str, Any]: 성공 시 (error, message, metadata)를 반환합니다. Metadata를 생략하고 저장했다면 빈 dictionary입니다. 실패하면 음수 error와 None payload를 반환합니다.

Usage

새 프로세스에서 학습을 재개하는 예시입니다.

error, message, trainer = PairedTrainer.build(config, data)
assert error >= 0, message
error, message, _ = trainer.to_device(device)
assert error >= 0, message
error, message, metadata = trainer.resume("training.pause.pt", password=password)
assert error >= 0, message
# pause() 예제에서 caller가 저장한 진행률을 사용합니다.
for step in range(metadata["completed_steps"], total_steps):
    error, message, result = trainer.train_one_step()
    assert error >= 0, message

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)

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