[python] heatmap - gradcam

모델의 activation visualization 부분을 분석해보자.

 

simple_grad_cam

def simple_grad_cam(features, classifier, target_class):
    """
    calculate gradient map.
    """
    features = nn.Parameter(features)

    logits = torch.matmul(features, classifier)
    
    logits[0, :, :, target_class].sum().backward()
    features_grad = features.grad[0].sum(0).sum(0).unsqueeze(0).unsqueeze(0)
    gramcam = F.relu(features_grad * features[0])
    gramcam = gramcam.sum(-1)
    gramcam = (gramcam - torch.min(gramcam)) / (torch.max(gramcam) - torch.min(gramcam))

    return gramcam

The simple_grad_cam function is used to calculate the gradient map for a given feature map. The gradient map highlights the regions of the feature map that are most relevant to the prediction of a particular target class. This is done by computing the gradients of the target class with respect to the feature map, and then weighting the feature map by these gradients. The resulting weighted feature map, or gradient map, highlights the regions of the feature map that are most important for predicting the target class. This is a useful technique in visualizing the attention mechanism of a neural network, and can be used to understand how the network makes its predictions. The gradient map can also be used to generate a heatmap image that highlights the important regions of the input image that contributed to the prediction.

 

def simple_grad_cam(features, classifier, target_class):

features, classifier, target_class를 인자로 받는다.

features = nn.Parameter(features)

텐서를 autograd하기 위해 nn.Parameters객체로 변환한다.

logits = torch.matmul(features, classifier)

features와 classifier를 곱하여 logit을 계산한다.

feature block이 layer1~layer4까지 있어서 이거에 대해 한번씩 출력하는 거 같고

feature은 각 layer의 차원 그러니까 layer1의 경우 [b,48,48,384]를 가진다.

classifier는 각 channel에 대한 클래스 개수만큼 즉 위의 경우 [384,200]이 된다.

이걸 곱해서 logits을 생성하는데 matmul수행 결과는 [1,48,48,200]

뭔가 fc 층과 비슷한 연산이 아닐까 하는 생각

logits[0, :, :, target_class].sum().backward()

여기서 타겟 클래스에 대한 logit의 그레디언트를 구한다.

아무래도 logit의 특징이 [b, 48,48,200] 이기 때문에 마지막 차원이 클래스 인덱스인 것 같다.

따라서 target_class를 인덱스로 받아 logits을 인덱싱하고 이를 합쳐서 loss 값을 구한 뒤 backward하는 것

Gradient map을 구하기 위해서 input feature map에 대한 output logit의 그레디언트를 계산해야 함.

이를 위해서 로짓의 합을 구하고 backward()하는 것.

features_grad = features.grad[0].sum(0).sum(0).unsqueeze(0).unsqueeze(0)

타겟 클래스에 대해 'features' 텐서의 그레디언트를 추출한다.

이후 feature map의 공간 차원을 따라 그레디언트를 합산하고 입력 텐서의 모양과 일치하도록 unsqueeze

그러니까 지금 features는 [1, 48, 48, 384]인데 처음에 [0]을 하면서 [48,48,384]가 되고

h,w에 대해 각각 sum을 계산했기 때문에 채널 수인 384만큼 남은 것

이후 unsqueeze를 두번 해서 [1,1,384]

gramcam = F.relu(features_grad * features[0])

이제 이 features_grad를 가중치로 해서 input features에 곱해준다.

이렇게 하면 gradient-weighted feature map 완성

그리고 ReLU 활성화맵을 거친다.

gradcam의 차원은 [48,48,384]

gramcam = gramcam.sum(-1)

channel dimension을 따라서 gradient-weighted feature map을 합산한다.

그러니까 이전의 gradcam은 

이렇게 하면 single channel map이 됨 

차원은 [48,48]

(뭔가 gradcam에서 avgpool같은 역할)

gramcam = (gramcam - torch.min(gramcam)) / (torch.max(gramcam) - torch.min(gramcam))

gradcam을 정규화

 

위의 과정을 통해 layer block에 대한 single gradient map을 완성하였다.

 

get_heat

def get_heat(model, img):

모델과 이미지를 인자로 받는다.

with torch.no_grad():
    outs = model.forward_backbone(img.unsqueeze(0))

이미지 텐서를 backbone에 통과시킨다.

이때 그레디언트를 계산하지 않기 위해 torch.no_grad() 

simple_grad_cam에선 classifier와 feature map 간의 gradient를 구하기 때문에 여기서는 disable해도 됨

features = []
for name in outs:
    features.append(outs[name][0])

feature map이 dictionary에 담겨있기 때문에 이때 feature(value)만 추출해서 하나씩 리스트에 담아줌

layer1의 경우 feature map 차원이 [1,2304,384]이기 때문에 [0]번째를 추출하면 [2304,384]

layer_weights = [8, 4, 2, 1]
heatmap = np.zeros([args.data_size, args.data_size, 3])

레이어 별로 가중치를 정의하고 heatmap을 초기화함

for i in range(len(features)):
    f = features[i]
    f = f.cpu()
    if len(f.size()) == 2:
        S = int(f.size(0) ** 0.5)
        f = f.view(S, S, -1)

각 feature map에 대해서 전처리 시행

2D feature map으로 shape을 변형함

그러니까 위에서 [2304,384]였으니 이를 [48,48,384]의 2D array로 변경

    gramcam = simple_grad_cam(f.unsqueeze(0), classifier=torch.ones(f.size(-1), 200)/f.size(-1), target_class=args.target_class)
    gramcam = gramcam.detach().numpy()
    gramcam = cv2.resize(gramcam, (args.data_size, args.data_size))

(이것도 for문 내부임)

근데 unsqueeze할거면 왜 굳이 인덱싱했지..?

암튼 [1,48,48,384]의 차원으로 다시 늘려주고

classifier는 [채널, 클래스 수]의 차원을 갖도록 초기화한다.

그리고 각각 1/(클래스 수)의 값을 갖도록 한다.

target_class는 0으로 고정되어 있는데 어떤 라벨을 써도 결과가 똑같았다.

그냥 gradient를 구하는게 목적이기 때문에 상관없나봄

그렇게 해서 simple_grad_cam을 통해 현재 특징맵에 대한 gradient map을 계산하고 input image size와 맞추기 위해 resize

	heatmap[:, :, 2] += layer_weights[i] * gramcam

gradcam의 값을 해당되는 레이어가중치와 곱해서 계속해서 heatmap에 더해준다.

이때 heatmap은 에서 선언하였듯 [img_size, img_size,3]의 차원을 갖는데

이때 3은 RGB 채널을 나타낸다.

즉, rgb중 2개는 0으로 채워져있고 하나만 값을 갖기 때문에 (여기서는 빨간색) 어텐션된 부분이 빨간색으로 나타날것

heatmap = heatmap / sum(layer_weights)
heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min())
heatmap[heatmap < args.threshold] = 0 # threshold
heatmap *= 255
heatmap = heatmap.astype(np.uint8)

heatmap을 정규화하고

threshold를 정해서 임계값보다 낮으면 0으로 처리한다.

이후 numpy array로 변환