Pytorch: 기능 요청: load_state_dict는 파일 이름을 가져와야 합니다.

에 만든 2017년 05월 31일 · 3코멘트 · 출처: pytorch/pytorch

높은 메모리 부족 상황에서 다음은 일반적으로 발생합니다.

모델 생성
체크포인트 파일에서 state_dict 읽기(GPU에 로드)
model.load_state_dict(들)

메모리 부족으로 인해 일반적인 해결 방법은 먼저 다음을 수행하는 것입니다.

s = torch.load('my_file.pt', map_location=lambda storage, loc: storage)

그런 다음 s 을 model 에 로드합니다.

이것은 우리가 피할 수 있어야 하는 매우 일반적인 시나리오이며 이 시나리오에는 몇 가지 함정이 있을 수 있습니다. 부분 GPU 부분 모델에서 일어나는 일, 다중 GPU 모델에서 일어나는 일...

load_state_dict가 파일 이름을 직접 가져오면 기존 매개 변수 저장소를 삭제하고 즉시 새 매개 변수로 설정할 수 있으므로 추가 메모리가 필요하지 않습니다.

feature nn triaged

출처

soumith

가장 유용한 댓글

load_state_dict 가 파일 이름을 사용하는 경우 map_location 매개변수도 허용해야 합니다. 나에게 일반적인 상황은 클러스터 시스템에 체크포인트를 저장한 다음 내 맥북에 로드하는 것입니다(따라서 CPU에 매개변수를 로드해야 함)

alykhantejani 에 2017년 06월 01일

👍3

모든 3 댓글

옵티마이저 state_dicts에도 동일하게 적용됩니다. Adagrad와 같은 일부 옵티마이저의 경우 체크포인트가 크며 동일한 메모리 압력 상황이 발생할 수 있습니다. 옵티마이저에는 .cuda() 조차 없으므로 먼저 state_dict를 CPU에 수동으로 로드한 다음 수동으로 일부를 GPU에 복사해야 합니다.

오늘 @aszlam 을 도우면서 이 문제를 만났습니다.

soumith 에 2017년 05월 31일

alykhantejani 에 2017년 06월 01일

👍3

나와 @szagoruyko 는 직렬화된 모델용 HDF5 형식의 팬입니다. 이 제안과 잘 어울릴 수 있다면

vadimkantorov 에 2017년 06월 06일

이 페이지가 도움이 되었나요?

0 / 5 - 0 등급

Pytorch: 기능 요청: load_state_dict는 파일 이름을 가져와야 합니다.

가장 유용한 댓글

모든 3 댓글

관련 문제