Skip to content

Resumable DataLoader

InfiniteRandomSampler와 MemorySafeDataLoader에 resumable=True를 지정하면 실제 소비한 sample 위치를 저장하고 새 프로세스에서 다음 batch부터 이어갈 수 있습니다. 두 클래스의 기본값은 False이며 기존 sampling과 iterator 동작을 유지합니다.

import torch

from SaigeToolkit.data.dataloader.memory_safe_dataloader import MemorySafeDataLoader
from SaigeToolkit.data.sampler.infinite_random_sampler import InfiniteRandomSampler

def make_loader(dataset):
    sampler = InfiniteRandomSampler(
        dataset, generator=torch.Generator().manual_seed(31), resumable=True,
    )
    return MemorySafeDataLoader(
        dataset, sampler=sampler, batch_size=8, num_workers=2,
        persistent_workers=True, resumable=True,
    )

loader = make_loader(dataset)
iterator = iter(loader)
batch = next(iterator)
# 이 batch의 optimizer step이 성공한 다음 Trainer checkpoint에 함께 보관합니다.
loader_state = loader.state_dict()
torch.save(loader_state, "loader-state.pt")
loader.close()

# 새 프로세스: 동일한 순서의 dataset을 준비한 뒤 sampler와 loader도 새로 만듭니다.
resumed_loader = make_loader(dataset)
resumed_loader.load_state_dict(torch.load("loader-state.pt", weights_only=False))
iterator = iter(resumed_loader)
next_batch = next(iterator)
resumed_loader.close()

Windows에서 worker를 사용한다면 위 실행 부분을 if __name__ == "__main__": 안에 둡니다. Dataset과 worker callback 정의는 worker가 import할 수 있는 module에 둡니다.

Sample 위치와 augmentation

worker의 prefetch 위치와 학습에 전달한 위치를 별도로 관리합니다. state_dict()는 전달 완료된 batch까지 기록합니다. load_state_dict()는 기존 worker와 미소비 prefetch를 정리하고 그 다음 sample부터 다시 생성합니다. Resumable 모드에서 iter(loader)를 여러 번 호출하면 현재 stream을 공유합니다. Shape 확인이나 logging을 위한 next(iter(loader))도 학습 batch를 소비하고 commit합니다. 이 호출로 읽은 batch는 기존 학습 iterator에서 다시 반환되지 않으므로, 관찰할 때는 학습 loop가 이미 받은 batch를 사용하세요. 소비하지 않는 별도 peek 기능은 제공하지 않습니다.

Batch 전달 중 fetch 또는 소비 위치 기록이 실패하면 해당 iterator는 더 사용할 수 없습니다. Fetch만 실패했다면 원인을 해결한 뒤 close()하고 iter(loader)로 iterator를 새로 만들어 마지막 소비 위치부터 다시 읽을 수 있습니다. Trainer가 보관한 iterator도 교체해야 합니다. 소비 위치 기록이 실패한 경우에는 close() 이후에도 새 iterator와 state_dict()를 차단합니다. 실패 전에 보관한 정상 loader checkpoint를 load_state_dict()로 복원해야 합니다.

close()는 worker를 종료하고 기존 iterator를 무효화합니다. Loader와 iterator를 모두 버리면 참조 순환 없이 worker가 정리되지만, iterator를 계속 보관하면 아직 사용 중인 stream으로 간주합니다. 정리 시점을 보장하려면 close()를 호출합니다.

각 sample의 augmentation seed는 loader state의 augmentation_seed와 stream 내 추출 순번에서 결정됩니다. 같은 이미지가 다시 뽑히면 다른 순번을 사용하므로 augmentation도 달라질 수 있습니다. 동일한 순번을 resume하면 같은 seed를 사용합니다. Loader 생성 시 Torch CPU RNG에서 seed를 한 번 뽑아 loader state에 저장하며, 복원 시 저장된 seed를 사용합니다. 별도의 생성자 옵션은 제공하지 않습니다.

전체 dataset[index] 실행을 Python random, NumPy global RNG, Torch CPU RNG scope로 감싸므로 그 안의 augmentation과 random crop을 함께 재현합니다. collate_fn도 batch별로 별도 seed를 사용합니다. Scope 종료 시 원래 RNG를 복원하여 num_workers=0에서도 모델이나 다른 loader의 RNG에 영향을 주지 않습니다. Dataset/collate scope는 CUDA RNG를 변경하지 않습니다.

지원 범위

  • 일반 Map-style dataset과 순서대로 전달하는 loader를 지원합니다. MapDataPipe와 IterableDataset의 resume는 지원하지 않습니다. InfiniteRandomSampler(resumable=True) 또는 Toolkit의 기본 batch sampler (naive, class_balance, balance)에 resumable=True를 지정할 수 있습니다. InfiniteRandomSampler의 subclass, 임의의 custom batch_sampler, in_order=False는 지원하지 않습니다. Subclass의 iterator와 저장 cursor가 달라질 수 있어 resume에서는 거절하며, sampler와 loader가 모두 resumable=False인 기존 경로에는 이 제한을 적용하지 않습니다.
  • Sampler와 loader의 resumable 설정을 함께 지정해야 합니다. Resumable sampler를 non-resumable MemorySafeDataLoader에 연결하면 ResumableModeRequiredError로 거절합니다. Sampler를 직접 순회하거나 일반 PyTorch DataLoader에 연결할 때는 실제 소비한 뒤 advance(count) 또는 advance_batch()를 호출해야 저장 cursor가 진행합니다.
  • Resumable loader는 하나의 활성 iterator를 공유합니다. next(iter(loader))를 shape 확인이나 logging용 peek로 사용하면 그 batch가 학습 stream에서 소비됩니다.
  • collate_fn은 loader 생성자의 인자로 전달해야 합니다. Resumable loader 생성 후 loader.collate_fn = ...로 교체하는 동작은 지원하지 않습니다. 내부 sample을 풀고 RNG를 복원하는 wrapper가 빠져 batch 구조와 augmentation 재현성이 깨질 수 있습니다. 복원할 새 loader에도 동일한 collate 설정을 생성자에 전달해야 합니다.
  • Loader마다 별도 sampler를 생성해야 합니다. 같은 sampler를 여러 loader에 연결하거나 활성 iterator가 있는 상태에서 다른 loader에 재사용하지 마세요. 이 조건을 별도 registry로 검사하지는 않습니다. 연결된 sampler의 advance()나 load_state_dict()를 외부에서 직접 호출하지 말고, 활성 loader의 저장·복원 API를 사용해야 합니다.
  • 동일한 dataset index 순서, 파일 내용, transform 설정과 라이브러리 버전이 필요합니다. Loader는 dataset 크기와 batch 설정을 검사하지만 파일 내용과 transform config의 동일성은 모듈에서 검사해야 합니다.
  • Python/NumPy/Torch CPU global RNG를 사용하는 augmentation을 지원합니다. RNG scope는 process마다 한 thread에서 실행해야 하며, DataLoader의 여러 worker process는 사용할 수 있습니다. 별도 np.random.Generator나 torch.Generator, stateful custom transform, 외부 서비스, 동시에 global RNG를 사용하는 thread는 별도 state 저장 또는 RNG 주입이 필요합니다.
  • 모듈이 추가로 보관하는 patch buffer, 여러 loader 사이의 선택 상태, 비동기 queue 순서, 학습 중 GPU augmentation은 해당 모듈에서 복원해야 합니다. BalancedSampler(calculate_patch_number=True)가 갱신하는 dataset.n_patch 같은 부가 metadata도 sampler checkpoint에 포함되지 않습니다. Loader 길이나 진행률에 사용한다면 모듈에서 별도의 복원 계약을 정해야 하며, sample stream 일치만으로 값의 일치를 보장하지 않습니다.
  • 학습 시작부터 resumable 모드를 선택해야 합니다. 기존 모드와 같은 seed를 주더라도 augmentation 순서가 달라질 수 있습니다. 초기 DET PR에서 사용한 구 loader state는 augmentation seed가 없으므로 이 형식으로 복원할 수 없습니다.

Trainer에서 저장할 상태

기존 batch sampler 사용

Class balancing을 사용하는 모듈은 기존 sampler를 그대로 확장합니다. repeat, shuffle, drop_last와 class/domain별 반복 비율을 유지하며, 마지막의 작은 batch도 실제 sample 수만큼 소비 위치를 진행합니다.

from SaigeToolkit.data.batch_sampler import build_batch_sampler

batch_sampler = build_batch_sampler(
    _target_="class_balance", dataset=dataset, batch_size=8,
    repeat=True, shuffle=True, drop_last=True, resumable=True,
)
loader = MemorySafeDataLoader(
    dataset, batch_sampler=batch_sampler, num_workers=2, resumable=True,
)

State에는 현재 cycle의 index batches와 위치, 누적 소비 수, 다음 cycle을 재생성할 seed를 저장합니다. Prefetch는 복사된 상태만 진행합니다. 복원 시 batch 크기뿐 아니라 class/domain balancing 설정과 cycle/offset/consumed의 일관성도 확인합니다.

같은 cycle의 batch 목록은 prefetch와 소비 위치 기록에서 재사용합니다. Cache는 worker의 prefetch 범위로 제한하며 소비가 끝난 cycle은 제거합니다. BalancedSampler의 n_patch 갱신 시점과 값은 유지하고, 외부에 반환하는 batch list는 복사해 내부 상태를 보호합니다. Cache는 checkpoint에 추가되지 않으며 복원 시 비웁니다.

오류 처리

Resume 검증 오류는 SaigeToolkit.error.error의 SaigeToolkitError 계열로 반환합니다. 원인별 Error 클래스에 고유 code와 고정 메시지를 정의하며, 호출부에서는 raise TrainingPausedError처럼 인자 없이 발생시킵니다. 사용하지 않는 code는 제거하되 남아 있는 code의 번호는 유지합니다. 제품의 error handler와 get_error_message(code)에서 같은 기본 메시지를 얻습니다.

대표적인 오류는 다음과 같습니다.

Error Code 의미
TrainingResumeCheckpointFormatError 321 pause checkpoint가 아님
TrainingPausedError 324 paused trainer에 학습 요청
ResumableBatchConfigurationError 338 비어 있는 dataset이나 유효한 batch를 생성할 수 없는 설정
TrainingResumeVersionMismatchError 329 module version 불일치
ResumableSamplerGeneratorError 351 sampler RNG state 손상

모든 클래스는 SaigeToolkitError만 상속합니다. Python 호출부에서는 구체 Error 타입으로 catch하고, 메시지는 str(error) 대신 error.message로 확인합니다. 기존 Error와 같이 API의 head와 Toolkit repo code를 붙이므로 324는 Toolkit API에서 -300324, Vision 모듈 API에서 -400324입니다. PyTorch worker가 추가하는 context와 원인 예외는 유지합니다. 미구현 adapter hook의 NotImplementedError, attribute protocol의 AttributeError, 검증 helper의 인자 오류(TypeError/ValueError)와 검증 실패(AssertionError), 외부에서 발생한 예외는 제품 Error code로 등록하지 않습니다.

모듈별 Trainer 연결

Loader state는 전체 학습 checkpoint의 한 부분입니다. Model, optimizer, scheduler, AMP scaler, 학습 step, module 고유 상태를 함께 저장하고 완료된 optimizer step에서만 checkpoint를 만들어야 합니다. 실패한 step의 loader 위치만 저장해서는 안전하게 재개할 수 없습니다.

SaigeToolkit.util.training_state는 다음 공통 utility를 제공합니다.

  • cpu_state(value): 중첩된 tensor state를 독립적인 CPU snapshot으로 복사합니다.
  • random_state(device): Python/NumPy/Torch CPU와 지정한 CUDA device의 RNG를 저장합니다.
  • restore_random_state(state, device): 저장한 main process RNG를 복원합니다.
  • restore_grad_scaler_state(scaler, state, optimizers): scaler state를 복원하고 실패한 AMP step의 임시 상태를 정리합니다. Fused optimizer에 남은 AMP 임시 속성도 제거합니다.

Trainer는 모델과 loader 및 iterator를 복원한 마지막에 main process RNG를 복원해야 합니다. Iterator 생성 자체가 난수를 소비할 수 있기 때문입니다. GPU kernel의 bitwise 결정성은 입력 augmentation의 재현과 별도로 검증해야 합니다.

SaigeToolkit.learning.training_resume.TrainingResumeMixin은 여러 모듈이 이 순서를 공유하도록 하는 opt-in helper입니다. 기존 BaseTrainer의 checkpoint 형식은 바꾸지 않습니다. 원본 config로 _initialize_training_resume()을 호출한 뒤 가장 바깥의 학습 step에 @training_step, state를 변경하는 다른 작업에 @serialized_training을 적용합니다. 각 모듈은 다음 hook을 구현합니다.

  • _resume_objects(): model, optimizer, scheduler, scaler 등 state_dict를 가진 객체.
  • _resume_loaders(): 저장할 training loader.
  • _resume_shared_keys(): 일반 checkpoint와 resume object가 같은 이름과 state를 사용하는 key 목록. 기본값은 빈 tuple이며, 변환한 export state는 제외합니다.
  • _resume_extra_state() / _load_resume_extra_state(): counters, GAN image pool, validation warm-up 여부 등 모듈 고유 상태.
  • _rebuild_resume_iterators(): loader 복원 후 iterator 생성.
  • _validate_resume(): 모듈별 지원 조건 검사.

일반 checkpoint와 resume objects에 같은 model을 저장한다면 다음처럼 선언할 수 있습니다.

def _resume_shared_keys(self) -> tuple[str, ...]:
    return ("model",)

pause()는 일반 checkpoint를 CPU로 복사한 뒤 해당 snapshot을 resume object에서도 참조합니다. 같은 model의 state_dict() 호출과 tensor 복사를 반복하지 않으며, 기존 key와 checkpoint 형식은 유지합니다. 이전에 별도 사본을 저장한 pause checkpoint도 읽을 수 있습니다. 같은 이름이라는 이유만으로 자동 공유하지는 않습니다. Export용으로 가공하거나 일부만 저장하는 state는 공유 목록에서 제외해야 합니다. 일반 checkpoint나 resume objects에 없는 key와 training_state 항목은 거절합니다. 이 최적화는 pause 경로에만 적용되며 일반 학습 step의 처리량 개선을 뜻하지 않습니다.

Mixin은 완료된 step에서만 저장하고 Windows에서도 writable handle로 fsync한 뒤 파일을 교체합니다. Restore 도중 실패한 Trainer는 계속 학습하지 못하도록 잠급니다. 학습 step의 실패는 재시도를 막지 않으며, 다음 step이 성공하거나 정상 checkpoint를 복원하기 전까지 pause()만 차단합니다. 재시도는 이미 변경된 weight·gradient·RNG나 소비한 batch를 되돌리지 않습니다. 실패 전 상태가 필요하면 resume()으로 복원해야 합니다. AMP나 데이터 오류는 원인을 해결하고 해당 객체의 상태도 정리한 뒤 재시도해야 합니다. 같은 module version/config/data 순서/device 종류/AMP 설정이 필요합니다. 공통 Mixin은 load 항목을 제외한 config의 모든 nested dictionary에서 key 삽입 순서까지 비교합니다. 값이 같더라도 JSON의 sort_keys=True 등으로 key 순서를 바꾸면 resume가 거절될 수 있으므로 원래 config의 순서를 보존해야 합니다. Config에 명시적인 torch.Generator를 전달한다면 첫 build 이전의 초기 상태로 재생성하거나 미리 복사한 config를 사용해야 합니다. 이미 사용하여 RNG 상태가 바뀐 Generator 객체를 같은 config에 넣어 재사용하면 signature가 달라져 resume가 거절될 수 있습니다. 실패한 AMP step 뒤 같은 Trainer에 checkpoint를 복원할 때도 scaler와 optimizer의 임시 상태를 정리하므로, 복원된 step 경계에서 다시 학습할 수 있습니다.

검증

모듈의 adapter는 SaigeToolkit.util.training_resume_test.assert_training_resume()으로 중단 없이 진행한 학습과 pause/resume한 학습을 비교할 수 있습니다.

from SaigeToolkit.util.training_resume_test import assert_training_resume

assert_training_resume(
    factory=make_trainer,  # 같은 초기 조건으로 매번 새 Trainer를 생성합니다.
    run_step=lambda trainer: trainer.train_one_step(),
    checkpoint_path="pause-test.pt",
    pause_after=3,
    total_steps=7,
    snapshot=lambda trainer: {"step": trainer.step},
)

run_step은 실행 시간처럼 매번 달라지는 값을 제외한 loss나 batch를 반환해야 합니다. snapshot은 저장 hook을 재사용하지 않고 실제 counter, crop buffer, loader 위치 등을 직접 읽어야 합니다. 그래야 저장 hook에서 빠뜨린 상태도 검출할 수 있습니다. 두 callback의 관찰 부분이 난수나 학습 상태를 변경해서는 안 됩니다.

Helper는 pause 전후와 resume 직후의 상태, 이후 step 결과와 최종 상태를 비교합니다. 새 Trainer를 생성한 뒤 RNG를 일부 진행시켜 constructor의 seed 초기화가 RNG 복원 누락을 가리지 않게 합니다. Paused step 차단도 검사합니다. 각 단계의 Trainer와 loader를 정리하고 호출자의 RNG를 복원합니다. 공통 helper는 CPU의 같은 process에서 실행되므로 별도 process의 spawn, GPU kernel 결정성이나 제품 executable의 재시작 테스트를 대신하지 않습니다.

CUDA_VISIBLE_DEVICES='' python -m pytest -q test/test_resumable_dataloader.py \
  test/test_resumable_batch_sampler.py test/test_training_resume_mixin.py \
  test/test_training_resume_contract.py test/test_training_state.py test/test_scheduler_resume.py

공통 테스트는 worker prefetch, 반복 sample의 seed, random collate, 실제 augmentation의 image/bbox, 새 프로세스와 spawn 재개, RNG scope의 예외 복원과 기존 모드 호환성을 검사합니다. 실제 loader와 Dropout/Adam/StepLR를 연결한 uninterrupted-vs-resume 비교 및 save/fsync/replace 실패 후 checkpoint 보존과 재시도도 검사합니다.

state_dict()의 반환 타입은 sampler의 SamplerState, loader의 LoaderState로 정의되어 있습니다. 저장된 format, cursor, index와 RNG state는 복원 전에 검사합니다. Config signature는 설정의 일치 여부만 확인하므로 손상된 state를 검출하는 이 검사를 대체하지 않습니다. 잘못된 loader state는 현재 worker와 iterator를 종료하기 전에 거절합니다. RNG helper도 RandomState를 반환하므로 편집기에서 checkpoint의 key와 값 타입을 확인할 수 있습니다. 실제 저장 형식은 기존 dictionary입니다.

state = loader.state_dict()  # LoaderState
consumed: int = state["sampler"]["consumed"]
augmentation_seed: int = state["augmentation_seed"]
loader.load_state_dict(state)

관련 모듈의 type hint는 Cython의 암묵적 타입 변환에 사용하지 않습니다. Python과 배포 extension에서 잘못된 flag, cursor, seed에 대한 기존 validation이 동일하게 동작하는지 regression test로 확인합니다. CI에서는 Python 구현과 Cython 배포 구현 모두에서 실행합니다.

inspect.signature(MemorySafeDataLoader)는 기존 DataLoader 생성자 인자와 새 keyword-only resume 옵션을 함께 노출합니다. typing.get_type_hints(MemorySafeDataLoader.__init__)로 인자 타입도 조회할 수 있으며, 별도 생성자를 정의한 subclass의 signature는 유지합니다.

cpu_state()는 모델 state_dict()의 OrderedDict와 _metadata를 복제하여 PyTorch가 모델 state를 load할 때 사용하는 version 정보도 보존합니다.