torch.load — 역직렬화·복원 API
torch.load — 역직렬화·복원 API
torch.load(f) 는 torch.save 로 저장한 객체를 파일에서 읽어 메모리로 되살리는 함수예요. pickle 역직렬화를 사용하고, 디바이스 매핑을 지정할 수 있어요.
사용법
checkpoint = torch.load('checkpoint.pt', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
map_location: 저장 당시 GPU 텐서를 로드 시 CPU나 특정 GPU로 옮길 때 쓴다(예:map_location='cpu', 람다로storage.cuda(dev)).weights_only=True(권장): 오직 텐서 및 기본 타입만 로드해, pickle 보안 위험을 줄인다.
주의
- 비신뢰 소스의 pickle은 인증되지 않은 코드를 실행할 수 있어 weights_only 사용을 권장한다.
- 로드 후에는 모델 구조를 먼저 만들고
load_state_dict로 상태를 넣는다. - 파일이 기기 간에 이동한 경우
map_location으로 텐서 디바이스를 조정한다.