torchsummary "Forward/backward pass size"

ResNeXt 모델 구현 중 구현이 잘되었는지 torchsummary를 사용해 파라미터 수를 사용해 확인하는 과정에서 이상한 점을 발견

• 파라미터 사이즈는 동일
• Forward/backward pass size 에서 45%가량 차이가 발생
• VRAM에서는 그 정도의 차이가 보이지 않음

직접 구현

================================================================  
 Total params: 25,028,904                                          
 Trainable params: 25,028,904                                      
 Non-trainable params: 0                                           
 ----------------------------------------------------------------  
 Input size (MB): 0.57                                             
 Forward/backward pass size (MB): 525.27                           
 Params size (MB): 95.48                                           
 Estimated Total Size (MB): 621.32                                 
 ----------------------------------------------------------------  

torchvision official Model

================================================================
Total params: 25,028,904
Trainable params: 25,028,904
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 361.78
Params size (MB): 95.48
Estimated Total Size (MB): 457.83
----------------------------------------------------------------

뭐지?

• params 개수가 동일
• 네트워크 구조가 같은데 (실수가 없었다면)
• Forward/backward pass size (MB)가 다르다??
• Estimated Total Size (MB)도 다르다?

무시하기엔 너무나도 큰 차이를 보인다.

실행해서 메모리 사용량과 학습 속도를 살펴보자

직접 구현

NVIDIA **** | 79'C,  45 %,  95 % | 10395 / 24267 MB 
train loop speed 4.89it/s

cat summary_self.txt | grep ReLU | wc -l 
49
cat summary_self.txt | grep Conv2d | wc -l
53
cat summary_self.txt | grep BatchNorm2d | wc -l
53

torchvision official Model

NVIDIA **** | 73'C,  40 %,  95 % | 10421 / 24267 MB 
train loop speed 4.89it/s

cat summary_official.txt | grep ReLU | wc -l 
49
cat summary_official.txt | grep Conv2d | wc -l
53
cat summary_official.txt | grep BatchNorm2d | wc -l
53

VRAM 사이즈에서 조금 차이가 나긴 하지만 신경 쓰이는 정도는 아닌 듯 하고
summary에서 확인되는 레이어 숫자도 같게 나오는데 왜!!
Identity 레이어가 들어가서 그런가?

확인해 보자

Identity Layer 실험

# 수정전
class ConvBlock(nn.Module):

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        **kwargs: Any,
    ):
        super().__init__()
        norm = kwargs.get('norm', True)
        act = kwargs.get('act', True)

        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kwargs.get('kernel_size'),
            stride=kwargs.get('stride', 1),
            padding=kwargs.get('padding', 0),
            groups=kwargs.get('groups', 1),
            bias=kwargs.get('bias', False),
        )
        self.norm = nn.BatchNorm2d(out_channels)
        self.act = nn.ReLU(inplace=True) if act else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv(x)
        x = self.norm(x)
        x = self.act(x)
        return x

수정후
...
        super().__init__()
        act = kwargs.get('act', True)
        layer = []

        layer += [
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kwargs.get('kernel_size'),
                stride=kwargs.get('stride', 1),
                padding=kwargs.get('padding', 0),
                groups=kwargs.get('groups', 1),
                bias=kwargs.get('bias', False),
            )
        ]
        layer += [nn.BatchNorm2d(out_channels)]
        if act:
            layer += [nn.ReLU(inplace=True)]

        self.block = nn.Sequential(*layer)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.block(x)

================================================================
 Total params: 25,028,904
 Trainable params: 25,028,904
 Non-trainable params: 0
 ----------------------------------------------------------------
 Input size (MB): 0.57
 Forward/backward pass size (MB): 471.67
 Params size (MB): 95.48
 Estimated Total Size (MB): 567.72
 ----------------------------------------------------------------
NVIDIA **** | 59'C,  31 %,  95 % | 10393 / 24267 MB |
train loop speed 4.96it/s

!!! 이번엔 또 조금 줄었다 !!!

Identity가 없어졌으니 없어진 만큼 메모리는 덜먹긴 할텐데 아직도 Forward/backward pass size에서 차이가 크게 나타난다!
이상..허다 ...

torchsummary 코드를 까보자 !

# torchsummary 

def summary 
	...
    for layer in summary:
        # input_shape, output_shape, trainable, nb_params
        line_new = "{:>20}  {:>25} {:>15}".format(
            layer,
            str(summary[layer]["output_shape"]),
            "{0:,}".format(summary[layer]["nb_params"]),
        )
		...
        ## 여기다! 
        total_output += np.prod(summary[layer]["output_shape"])
        ...
        print(line_new)
	
    # assume 4 bytes/number (float on cuda).
    total_output_size = abs(2. * total_output * 4. / (1024 ** 2.))
	...
    print("Forward/backward pass size (MB): %0.2f" % total_output_size)
    ...
    print("----------------------------------------------------------------")

알았다.

• Forward/backward pass size는 출력되는 layer의 출력 shape 곱의 합산
• summary에서 만들어지는 한줄 한줄을 바로 합산하기 때문에 nn.Module로 묶어놓거나 Identity 레이어가 있으면 해당 레이어가 summary로 출력되면서 그만큼의 shape만큼 사이즈가 늘어나게 측정이 누적됨

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
	   		Conv2d-1         [-1, 64, 112, 112]           9,408
       BatchNorm2d-2         [-1, 64, 112, 112]             128
              ReLU-3         [-1, 64, 112, 112]               0
         ConvBlock-4         [-1, 64, 112, 112]               0 -> 이 부분이 원인
         ...
================================================================
Total params: 25,028,904
Trainable params: 25,028,904
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 471.67
Params size (MB): 95.48
Estimated Total Size (MB): 567.72
----------------------------------------------------------------

이해한 내용이 확실한지 확인해보자

input_size = (3, 128, 128)

# conv_block_w (수정전)
torchsummary.summary(
    conv_block_w,
    input_size=input_size,
    device='cpu',
)
# conv_block_wo (수정후)
torchsummary.summary(
    conv_block_wo,
    input_size=input_size,
    device='cpu',
)

>>> 
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 64, 64, 64]           3,072
       BatchNorm2d-2           [-1, 64, 64, 64]             128
          Identity-3           [-1, 64, 64, 64]               0
================================================================
Total params: 3,200
Trainable params: 3,200
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.19
Forward/backward pass size (MB): 6.00
Params size (MB): 0.01
Estimated Total Size (MB): 6.20
----------------------------------------------------------------
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 64, 64, 64]           3,072
       BatchNorm2d-2           [-1, 64, 64, 64]             128
================================================================
Total params: 3,200
Trainable params: 3,200
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.19
Forward/backward pass size (MB): 4.00
Params size (MB): 0.01
Estimated Total Size (MB): 4.20
----------------------------------------------------------------

torchsummary에서 Forward/backward pass size는 아래와 같이 구해진다.
total_output_size = abs(2. * total_output * 4. / (1024 ** 2.))

각 레이어 출력이 [64, 64, 64] 였으니

${each \ layer} = 2 * 64^3 * 4 / 1024^2$

정확하게 2MB 차이가 나도록 출력이 된다.

학습 forward/backward step에서 이게 어느 정도의 성능 하락을 일으키는지는 좀 더 알아봐야 하지만 실행시켰을 때 발생하는 성능 하락은 크게 보이지 않았기 때문에 성능 하락 관련 내용은 다음에 추가해 보려 한다.