본문 바로가기
DL|ML

torch의 model summary에서는 cuda가 필요한가요?

by 이든Eden 2021. 1. 5.

이 글은 이곳을 번역한 글이며 추가로 제 경험도 들어있습니다.

 

케라스에는 아주 편리하게 model.summary() 하면 레이어에 따른 파라미터 갯수가 쭉쭉 나온다.

파이토치는 따로 import 시켜야 그걸 할 수 있는데 아래 처럼 사용하면 된다.

from torchsummary import summary
summary(vgg, (3, 224, 224))

그러면 바로 아래 사진처럼 쭉쭉 나와야한다.

 

 

내가 참고한 글부터 이야기하자면 이렇다.

summary를 사용하면 아래와 같은 에러가 뜬다고 한다.

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

"input타입은 torch.cuda.FloatTensor고 weight타입은 torch.FloatTensor에요. 똑같이 맞춰주세요!" 

 

타입문제이지만 cuda를 사용한 타입인지 cpu를 사용한 타입인지의 문제다. gpu냐 cpu냐.. 내 문제는 타입은 아니지만 역시 gpu냐 cpu냐 하는 문제다. 내 문제는 아래와 같다.

 

torch hub에서 받은 Unet 레이어 상태 좀 보려고 근데 GPU를 쓰고있어서 그런지 RuntimeError: CUDA out of memory 가 나왔고 뭐지 싶었다. summary에서 gpu, cpu를 설정해줄 수 있는지 찾아보니 summary에는 device를 설정하는 매개변수인 device가 있다고 한다(말이 좀 이상한데 변수명도 device라는 이야기). 따라서 아래처럼 사용하면 된다.

 

from torchsummary import summary
summary(vgg, (3, 224, 224), device='cpu')

위와 같이 device를 cpu로 설정해주면 gpu를 사용하지 않기 때문에 원하던대로 레이어에 따른 파라미터 개수, output 모양이 나온다.