train loss = NaN되는 현상

[jsh0551]

mosaic으로 offline augmentation한 데이터로 knet+upernet 학습시킬 때 loss가 NaN으로 나오면서 학습이 터지는 현상이 있었습니다. 찾아보니 CNN+activation을 거치면 gradient가 급변하게 되는데 이때 loss 값이 크게 튀어서 나오는 현상이라고 합니다.

해결법은 모든 activation function 뒤에 batch normalization을 해주면 gradient 값이 안정화된다고 하네요. knet+upernet 같은 경우에는 decoder 학습 과정에서 이 현상이 발생했고, kernel_update_head 부분을 수정해서 해결했습니다.

kernel_update_head=[
            dict(
                type='KernelUpdateHead',
                num_classes=11,
                num_ffn_fcs=2,
                num_heads=8,
                num_mask_fcs=1,
                feedforward_channels=2048,
                in_channels=512,
                out_channels=512,
                dropout=0.0,
                conv_kernel_size=conv_kernel_size,
                ffn_act_cfg=dict(type='ReLU', inplace=True),
                with_ffn=True,
                feat_transform_cfg=dict(
                    conv_cfg=dict(type='Conv2d'), norm_cfg=norm_cfg), ## act_cfg=None을 norm_cfg=norm_cfg로 수정
                kernel_updator_cfg=dict(
                    type='KernelUpdator',
                    in_channels=256,
                    feat_channels=256,
                    out_channels=256,
                    act_cfg=dict(type='ReLU', inplace=True),
                    norm_cfg=dict(type='LN'))) for _ in range(num_stages)
        ]

[jsh0551]

에폭 20이 넘어가면서 NaN이 발생했습니다. optimizer_config의 타입으로 Fp16OptimizerHook설정을 했는데 float16 type으로 계산하기 때문에 속도는 빨라지는 대신 소실되는 값으로 인한 grad나 loss에서 NaN이 발생하는 경우가 많아지는 것 같습니다.

이를 방지하기 위한 장치로 loss_scale을 조절하여 방지할 수 있다고 했는데, 여러 값으로 실험했지만 계속 NaN이 떴습니다. 그래서 오래걸리지만 Fp16OptimizerHook을 쓰지 않고 학습 진행중입니다.