fvcore.nn.FlopCountAnalysis 모델 FLOPs 계산시 RuntimeError: Cannot insert a Tensor that requires grad as a constant. 에러

딥러닝 모델의 FLOPs를 계산할 때 fvcore.nn.FlopCountAnalysis 라이브러리를 주로 사용한다. 사용 예시는 아래와 같다.

$ from fvcore.nn import FlopCountAnalysis
$ flops = FlopCountAnalysis(model, input)
$ flops.total()
274656
$ flops.by_operator()
Counter({'conv': 194616, 'addmm': 80040})
$ flops.by_module()
Counter({'': 274656, 'conv1': 48600,
         'conv2': 146016, 'fc1': 69120,
         'fc2': 10080, 'fc3': 840})
$ flops.by_module_and_operator()
{'': Counter({'conv': 194616, 'addmm': 80040}),
 'conv1': Counter({'conv': 48600}),
 'conv2': Counter({'conv': 146016}),
 'fc1': Counter({'addmm': 69120}),
 'fc2': Counter({'addmm': 10080}),
 'fc3': Counter({'addmm': 840})}

출처 : https://github.com/facebookresearch/fvcore/blob/main/docs/flop_count.md

만약 모델 input이 2개 이상이라면?

flops = FlopCountAnalysis(model, (input1, input2))

이렇게 넣어주면 된다.

그리고 flops.total(), flop_count_table(flops) 등등 원하는 출력을 뽑아내면 된다. 개인적으로 flop_count_table(flops)이 깔끔하고 좋았다.

그런데 결과 출력 과정에서 간혹 이런 에러가 나오는 경우가 있다.

RuntimeError: Cannot insert a Tensor that requires grad as a constant.

이건 Multi GPU로 모델을 학습시켰을 경우에 발생한다. 자세히는 모르겠지만 DataParallel 했을 때 모델 weight 이름에 module이 붙는데 그게 문제가 되는 것 같다.

if isinstance(model, (nn.parallel.distributed.DistributedDataParallel, nn.DataParallel)):
	model = model.module

이렇게 하면 문제가 해결된다.