Yolov5에서 ModelEMA와 model fuse가 의미하는 것

 

 

Fuse란?

---

Yolov5는 대표적인 객체 탐지(Object detection) 모델 중 하나이다.
탐지(추론)를 위해 모델을 불러올 때, 그 코드를 살펴보면 attempt_load라는 함수를 이용한다.
* detect.py > DetectMultiBackend() > attempt_load()에서 찾을 수 있다.

그런데 이 함수를 잘 살펴보면 ckpt.get('ema'), 즉 모델의 'ema'라는 속성을 가져온다거나,
ckpt.fuse().eval() if fuse와 같이 불러온 모델을 fuse하도록 하는 것을 알 수 있다.

# yolov5/model/experimental.py

def attempt_load(weights, device=None, inplace=True, fuse=True):
    from yolov5.models.yolo import Detect, Model

    # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
    model = Ensemble()
    for w in weights if isinstance(weights, list) else [weights]:
        ckpt = torch.load(attempt_download(w), map_location=device)
        ckpt = (ckpt.get('ema') or ckpt['model']).float()  # FP32 model
        model.append(ckpt.fuse().eval() if fuse else ckpt.eval())  # fused or un-fused model in eval mode

    # 중략

    if len(model) == 1:
        return model[-1]  # return model

    # 생략

먼저 Fuse라는 단어의 사전적 의미를 잠깐 생각해보자.
"Join or blend to form a single entity", 즉 서로 다른 여러 가지를 섞고, 결합/융합시켜서 하나로 만든다는 뜻이다.
그렇다면 과연 무엇을 결합해서 하나로 만들고 싶은 걸까?

모델의 fuse() 함수를 찾아서 살펴보면 다음과 같은 내용이 있다.
주석의 내용과 함께 이해해보면, 모델의 각 모듈 중에서 Convolutional layer나 Depth-wise Convolutional layer가 있고,
동시에 해당 모듈에 'bn' 즉, Batch Normalization(배치 정규화)를 뜻하는 속성이 있으면 기존 Conv layer를 fuse_conv_and_bn이라는 함수를 통해 업데이트한다.
또한 기존 forward 메서드 대신, forward_fuse라고 하는 메서드로 대체한다.

# yolov5/models/yolo.py

class Model(nn.Module):
    # 중략

    def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
        for m in self.model.modules():
            if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
                m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
                delattr(m, 'bn')  # remove batchnorm
                m.forward = m.forward_fuse  # update forward
        self.info()
        return self

    # 생략

그러니까 결국 fuse()는 기존 Conv2d layer를 BatchNorm2d layer와 결합해 하나로 만드는 역할을 한다.
하지만 이런 과정이 왜 필요한 것일까?

 

 

Inference를 위한 Batch Normalization

---

Conv2d layer와 BatchNorm2d layer의 fusion을 이해하기 위해 Batch Normalization에 대해 깊게 이해할 필요가 있다.
Batch Normalization(줄여서 BN)이란 은닉층(Hidden layer)들에도 입력 데이터에서처럼 정규화 과정이 필요하다는 가정에서 시작되었다.

이를 처음 제시한 논문에서는 은닉층이 깊게 쌓이면 쌓일수록 각 Internal Covariate Shift(줄여서 ICS) 문제가 심각해질 것으로 예상했기 때문이다.
ICS가 심할수록 학습 속도가 느려지는 것은 물론, 학습 자체가 잘 이뤄지지 않아 높은 정확도를 얻기 어렵다.

$$BN(X) = \gamma \left ( \frac{X - \mu_{BN} }{\sigma_{BN}} \right ) + \beta$$

결과적으로는 해당 의도와 다르게 ICS 문제를 해결하지 않는 것으로 다른 논문을 통해 검증되었지만, 그와 별개로 성능 향상 자체에 크게 기여하고 있으며, 때문에 BN을 적용하는 것은 이제 당연시되고 있다.
* 자세한 내용은 이 영상을 참고하길 바란다.

https://i.ytimg.com/vi/Cb8NX3HiS4U/maxresdefault.jpg

우리가 주목해야 할 것은 BN을 언제, 어떻게 적용하느냐 이다.
모델 학습 과정에서는 BN을 적용하기 위한 파라미터들도 함께 학습하면서 최적의 값을 찾아간다.
하지만 추론 과정에서는 기존에 학습했던 파라미터 값들을 모델 파일에 함께 저장해뒀다가, 이를 불러와서 사용해야 한다.
이때 등장하는 것이 EMA(Exponential Moving Average, 지수 이동 평균)이다.

지수 이동 평균은 다른 말로 Exponentially Weighted Moving Average라고도 한다. 즉, 최근 데이터에 더 가중치를 주어 계산한 평균이라는 뜻이다.
'지수'라는 표현이 붙은 이유는 오래된 데이터부터 최근 데이터까지 적용되는 가중치가 지수적으로 변하기 때문이다.
수식으로 표현하면 다음과 같다.

$$V_t = \beta V_{t-1} + (1-\beta) \theta_{t}$$

* $V_t$: 현재 값, $V_{t-1}$: 이전 값, $\beta$: 모멘텀, $\theta_t$: 현재 데이터에 대한 bias

위 식을 전개해서 계산해보면 오래된 데이터일수록 0에 지수적으로 가까워지는 값을 얻게 됨을 알 수 있다.
* 지수 이동 평균에 대한 직관적 이해는 이 영상을 참고하길 바란다. 주식에서 지수 이동 평균을 사용하는 이유에 대해 예를 들어 설명하고 있다.

---

중간 정리를 해보자.
Yolov5 모델은 추론을 하기 위해 모델을 로드하는 과정에서 기존에 저장해 둔 모델의 'ema' 속성을 가져오고,
해당 값들을 이용해 Conv layer와 BN layer를 결합하는 fusion을 수행한다.

# yolov5/utils/torch_utils.py

class ModelEMA:
    def __init__(self, model, decay=0.9999, tau=2000, updates=0):
        # Create EMA
        self.ema = deepcopy(de_parallel(model)).eval()  # FP32 EMA
        self.updates = updates  # number of EMA updates
        self.decay = lambda x: decay * (1 - math.exp(-x / tau))  # decay exponential ramp
        for p in self.ema.parameters():
            p.requires_grad_(False)

    def update(self, model):
        # Update EMA parameters
        self.updates += 1
        d = self.decay(self.updates)

        msd = de_parallel(model).state_dict()  # model state_dict
        for k, v in self.ema.state_dict().items():
            if v.dtype.is_floating_point:  # true for FP16 and FP32
                v *= d
                v += (1 - d) * msd[k].detach()

    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
        # Update EMA attributes
        copy_attr(self.ema, model, include, exclude)

 

 

실제 Yolov5 코드로 살펴보기

---

이제 이를 실제로 학습 시에 어떻게 적용하는지 살펴보자.
train.py의 train 함수 코드를 보면, 우선 GPU 활용이 가능할 경우 ModelEMA를 생성하는 것으로 시작한다.
이전에 학습 도중 중단되었다면, 재시작(resume)을 위해 기존에 저장해 둔 checkpoint(줄여서 ckpt) 파일에서 'ema' 속성을 가져와 로드한다.

또한 실제로 학습 중 backward 이후 optimization을 수행하면서 ema.update()를 수행한다.

# yolov5/train.py

def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictionary
    # 중략

    # EMA
    ema = ModelEMA(model) if RANK in {-1, 0} else None

    # Resume
    start_epoch, best_fitness = 0, 0.0
    if pretrained:
        # Optimizer
        if ckpt['optimizer'] is not None:
            optimizer.load_state_dict(ckpt['optimizer'])
            best_fitness = ckpt['best_fitness']

        # EMA
        if ema and ckpt.get('ema'):
            ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
            ema.updates = ckpt['updates']

    # 중략

    for epoch in range(start_epoch, epochs):  # epoch -------
        # 중략

        # Optimize
        if ni - last_opt_step >= accumulate:
            scaler.step(optimizer)  # optimizer.step
            scaler.update()
            optimizer.zero_grad()
            if ema:
                ema.update(model)
            last_opt_step = ni

    # 생략

 

다음으로 BN이 언제 적용되는지 살펴보자.
앞서 살펴본 것처럼 DetectMultiBackend에서는 attempt_load()를 통해 모델을 가져와서 inference를 시작한다.

# yolov5/models/common.py

class DetectMultiBackend(nn.Module):
    # YOLOv5 MultiBackend class for python inference on various backends
    from models.experimental import attempt_download, attempt_load  # scoped to avoid circular import

    # 중략
    if pt:  # PyTorch
        model = attempt_load(weights if isinstance(weights, list) else w, device=device, inplace=True, fuse=fuse)
        stride = max(int(model.stride.max()), 32)  # model stride
        names = model.module.names if hasattr(model, 'module') else model.names  # get class names
        model.half() if fp16 else model.float()
        self.model = model  # explicitly assign for to(), cpu(), cuda(), half()
    # 중략

def forward(self, im, augment=False, visualize=False):
    # YOLOv5 MultiBackend inference
    # 중략

    if self.pt:  # PyTorch
        y = self.model(im, augment=augment, visualize=visualize) if augment or visualize else self.model(im)
    # 생략

이렇게 불러온 모델은 여러 모듈로 구성되는데, 그 중 우리가 살펴볼 것은 Conv 모듈이다.
앞서 Conv 또는 DWConv layer에 bn attribute이 존재할 경우에 Conv + BN fusion을 수행한다고 했다.
아래 코드에서 볼 수 있듯 DWConv layer에는 bn 속성이 없으므로 Conv 모듈만 살펴보도록 하자.

# yolov5/models/common.py

class Conv(nn.Module):
    # Standard convolution
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def forward_fuse(self, x):
        return self.act(self.conv(x))


class DWConv(Conv):
    # Depth-wise convolution class
    def __init__(self, c1, c2, k=1, s=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), act=act)

Conv 모듈을 살펴보면, 일반적인 nn.Conv2d와 nn.BatchNorm2d, 그리고 activation으로 구성되어 있는 것을 확인할 수 있다.
또한 일반적인 forward 메서드와 forward_fuse 메서드가 각각 따로 정의된다.
forward 메서드에서는 self.act(self.bn(self.conv(x))), 즉 일반적인 순서를 거쳐 각 연산을 순차적으로 수행한다.

반면에 forward_fuse 메서드에서는 self.act(self.conv(x))로 BN을 생략하는데,
그 이유는 앞서 살펴본 것처럼 model.fuse()를 수행하면서 이들의 결합, 즉 fusion을 위한 연산을 따로 수행할 것이기 때문이다.

# yolov5/models/yolo.py

class Model(nn.Module):
    # 중략

    def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
        for m in self.model.modules():
            if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
                m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
                delattr(m, 'bn')  # remove batchnorm
                m.forward = m.forward_fuse  # update forward
        self.info()
        return self

    # 생략

 

---

위 코드에서 볼 수 있듯이 m.conv = fuse_conv_and_bn(m.conv, m.bn)을 수행함으로써 fuse_conv_and_bn 함수의 return 값으로 기존 conv 모듈을 대체한다.
이후 forward 또한 forward_fuse로 대체하여, BN이 중복하여 수행되지 않도록 처리한다.

fuse_conv_and_bn 함수를 자세히 살펴보면, 'Prepare filters'와 'Prepare sptial bias'를 수행한다.
이들은 각각 Conv + BN의 결합된 형태의 weight와 bias를 연산하기 위한 과정을 나타낸다.

# yolov5/utils/torch_utils.py

def fuse_conv_and_bn(conv, bn):
    # Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
    fusedconv = nn.Conv2d(conv.in_channels,
                          conv.out_channels,
                          kernel_size=conv.kernel_size,
                          stride=conv.stride,
                          padding=conv.padding,
                          groups=conv.groups,
                          bias=True).requires_grad_(False).to(conv.weight.device)

    # Prepare filters
    w_conv = conv.weight.clone().view(conv.out_channels, -1)
    w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
    fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))

    # Prepare spatial bias
    b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
    b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
    fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reswhape(-1) + b_bn)

    return fusedconv

이를 설명하기 위해, 각 feature map($ C \times H \times W$)에 대한 Conv + BN fusion을 수식으로 나타내면 다음과 같다.

$$\hat{f_{i,j}} = W_{BN}\, \cdot\, (W_{conv}\cdot f_{i,j} + b_{conv}) + b_{BN}$$

* 각 기호에 대한 자세한 설명은 여기를 참고하길 바란다.

결과적으로 해당 feature map(= filter)에 대한 weights는 $W = W_{BN} \dot W_{conv}$,
bias는 $b = W_{BN} \dot b_{conv} + b_{BN}$이라고 할 수 있다.
이는 위 코드에서 말하는 filters와 spatial bias를 연산하는 기본 식이 된다.

 

마지막으로 Conv + BN fusion을 수행하기 전, 후 모델을 출력해 비교하면 다음과 같다.
모든 Conv 모듈 안에 포함된 BN layer가 Conv2d + BN fusion layer로 대체되었음을 알 수 있다.

# Before Fusion

Yolov5Detector(
  (model): Model(
    (model): Sequential(
      (0): Conv(
        (conv): Conv2d(3, 80, kernel_size=(6, 6), stride=(2, 2), padding=(2, 2), bias=False)
        (bn): BatchNorm2d(80, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)
        (act): SiLU(inplace=True)
      )
      (1): Conv(
        (conv): Conv2d(80, 160, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn): BatchNorm2d(160, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)
        (act): SiLU(inplace=True)
      )
      (2): C3(
        (cv1): Conv(
          (conv): Conv2d(160, 80, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (bn): BatchNorm2d(80, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)        
          (act): SiLU(inplace=True)
        )
        (cv2): Conv(
          (conv): Conv2d(160, 80, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (bn): BatchNorm2d(80, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)        
          (act): SiLU(inplace=True)
        )
        (cv3): Conv(
          (conv): Conv2d(160, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (bn): BatchNorm2d(160, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)       
          (act): SiLU(inplace=True)
        )
	# 이하 생략

# After Fusion

Yolov5Detector(
  (model): Model(
    (model): Sequential(
      (0): Conv(
        (conv): Conv2d(3, 80, kernel_size=(6, 6), stride=(2, 2), padding=(2, 2))
        (act): SiLU(inplace=True)
      )
      (1): Conv(
        (conv): Conv2d(80, 160, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        (act): SiLU(inplace=True)
      )
      (2): C3(
        (cv1): Conv(
          (conv): Conv2d(160, 80, kernel_size=(1, 1), stride=(1, 1))
          (act): SiLU(inplace=True)
        )
        (cv2): Conv(
          (conv): Conv2d(160, 80, kernel_size=(1, 1), stride=(1, 1))
          (act): SiLU(inplace=True)
        )
        (cv3): Conv(
          (conv): Conv2d(160, 160, kernel_size=(1, 1), stride=(1, 1))
          (act): SiLU(inplace=True)
        )
	
    # 이하 생략