직렬화

직렬화 (Serialization)

학습한 모델의 가중치를 파일로 저장해 두고 나중에 다시 불러오거나, 다른 환경으로 옮기고 싶을 때가 있죠. PyTorch의 torch.save() 와 torch.load() 는 텐서와 모듈 상태를 파이썬에서 쉽게 저장·로드하게 해 주고, C++ 쪽으로 직렬화하는 방법까지 다루는 개념이에요.

출처: PyTorch 공식 문서 — Serialization semantics

텐서 저장하고 불러오기

torch.save() 와 torch.load() 가 그 역할을 해요. 관례로 파일 확장자는 .pt 나 .pth 를 써요. 기본적으로 Python의 pickle을 사용하기 때문에 튜플·리스트·딕셔너리처럼 텐서를 품은 객체 통째로도 저장할 수 있어요.

>>> t = torch.tensor([1., 2.])
>>> torch.save(t, 'tensor.pt')
>>> torch.load('tensor.pt')
tensor([1., 2.])

>>> d = {'a': torch.tensor([1., 2.]), 'b': torch.tensor([3., 4.])}
>>> torch.save(d, 'tensor_dict.pt')
>>> torch.load('tensor_dict.pt')
{'a': tensor([1., 2.]), 'b': tensor([3., 4.])}

뷰 관계 보존

저장할 때 텐서의 뷰(view) 관계도 함께 보존돼요. 같은 저장소(storage)를 공유하는 텐서들은 로드한 뒤에도 그 관계를 유지해요. PyTorch는 저장소 객체와 텐서 메타데이터를 분리해서 저장하는데, 그래야 뷰 관계를 쉽게 재구성하고 파일 크기도 아낄 수 있어요.

>>> numbers = torch.arange(1, 10)
>>> evens = numbers[1::2]
>>> torch.save([numbers, evens], 'tensors.pt')
>>> loaded_numbers, loaded_evens = torch.load('tensors.pt')
>>> loaded_evens *= 2
>>> loaded_numbers
tensor([ 1,  4,  3,  8,  5, 12,  7, 16,  9])

다만 뷰를 저장하면 저장소 전체가 파일에 담겨 파일이 필요 이상으로 커질 수 있어요. 저장소보다 작은 텐서를 저장할 때는 clone() 해서 새 저장소(값만 담김)를 만들면 크기를 줄일 수 있지만, 뷰 관계는 사라져요.

>>> large = torch.arange(1, 1000)
>>> small = large[0:5]
>>> torch.save(small, 'small.pt')
>>> torch.load('small.pt').storage().size()
999

>>> torch.save(small.clone(), 'small.pt')
>>> torch.load('small.pt').storage().size()
5

모듈은 state dict로 저장하기

모듈을 통째로 저장하는 대신 state dict 만 저장하는 걸 권장해요. state dict 는 모든 파라미터와 영속 버퍼를 담고 있고, load_state_dict() 로 그 상태를 복원할 수 있어요. 호환성 면에서도 훨씬 안전해요.

>>> bn = torch.nn.BatchNorm1d(3, track_running_stats=True)
>>> bn.state_dict()
OrderedDict([('weight', tensor([1., 1., 1.])),
             ('bias', tensor([0., 0., 0.])),
             ('running_mean', tensor([0., 0., 0.])),
             ('running_var', tensor([1., 1., 1.])),
             ('num_batches_tracked', tensor(0))])

>>> torch.save(bn.state_dict(), 'bn.pt')
>>> new_bn = torch.nn.BatchNorm1d(3, track_running_stats=True)
>>> new_bn.load_state_dict(torch.load('bn.pt'))
<All keys matched successfully>

저장 파일 형식

PyTorch 1.6.0 이후 torch.save 는 기본적으로 압축하지 않은 ZIP64 아카이브를 만들어요. 아카이브 안에는 data.pkl(pickle 결과), byteorder(2.1.0부터), data/(저장소 개별 파일들), version 이 순서대로 들어 있어요. 각 파일의 로컬 헤더는 64바이트 배수로 정렬돼요.

`weights_only=True` 와 보안

버전 2.6부터 pickle_module 을 넘기지 않으면 torch.load 는 weights_only=True 를 사용해요. 이 옵션은 역직렬화 중 동적 import 를 막아 원격 코드 실행 공격의 표면을 줄여 주지만, 서비스 거부(DoS) 공격까지 막지는 못해요. 신뢰할 수 없는 체크포인트를 로드할 때는 허용 목록(allowlist)을 꼭 명시적으로 관리해야 해요.

직렬화