CNN은 어떻게 특정 패턴을 찾을까? PyTorch로 이해하는 Convolution과 Kernel

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np

# ==========================================
# 1. 입력 이미지 생성 (21x21)
# ==========================================

image = torch.zeros((21, 21), dtype=torch.float32)

# ==========================================
# 2. 3x3 X 패턴
# ==========================================

x3 = torch.tensor([
    [1,0,1],
    [0,1,0],
    [1,0,1]
], dtype=torch.float32)

image[1:4, 1:4] = x3

# ==========================================
# 3. 5x5 X 패턴
# ==========================================

x5 = torch.tensor([
    [1,0,0,0,1],
    [0,1,0,1,0],
    [0,0,1,0,0],
    [0,1,0,1,0],
    [1,0,0,0,1]
], dtype=torch.float32)

image[6:11, 6:11] = x5

# ==========================================
# 4. 7x7 X 패턴
# ==========================================

x7 = torch.tensor([
    [1,0,0,0,0,0,1],
    [0,1,0,0,0,1,0],
    [0,0,1,0,1,0,0],
    [0,0,0,1,0,0,0],
    [0,0,1,0,1,0,0],
    [0,1,0,0,0,1,0],
    [1,0,0,0,0,0,1]
], dtype=torch.float32)

image[13:20, 13:20] = x7

# ==========================================
# CNN 입력 형태
# ==========================================

image_tensor = image.unsqueeze(0).unsqueeze(0)

# ==========================================
# 5. Kernel 생성
# ==========================================

kernel3 = x3.unsqueeze(0).unsqueeze(0)
kernel5 = x5.unsqueeze(0).unsqueeze(0)
kernel7 = x7.unsqueeze(0).unsqueeze(0)

# ==========================================
# 6. Convolution 수행
# ==========================================

feature3 = F.conv2d(image_tensor, kernel3)

feature5 = F.conv2d(image_tensor, kernel5)

feature7 = F.conv2d(image_tensor, kernel7)

# ==========================================
# 7. Feature Map 출력
# ==========================================

print("\n==============================")
print("3x3 Kernel Feature Map")
print("==============================")
print(feature3.squeeze())

print("\n==============================")
print("5x5 Kernel Feature Map")
print("==============================")
print(feature5.squeeze())

print("\n==============================")
print("7x7 Kernel Feature Map")
print("==============================")
print(feature7.squeeze())

# ==========================================
# 8. 최대 응답값 및 위치
# ==========================================

for name, feature in [
    ("3x3", feature3),
    ("5x5", feature5),
    ("7x7", feature7)
]:

    fmap = feature.squeeze()

    max_value = fmap.max()

    max_pos = torch.nonzero(fmap == max_value)

    print(f"\n{name} Kernel")
    print("최대 응답값 =", max_value.item())
    print("최대 응답 위치 =")
    print(max_pos)

# ==========================================
# 9. 시각화
# ==========================================

fig, axes = plt.subplots(
    2,
    2,
    figsize=(14,10)
)

# ------------------------------------------
# Input Image
# ------------------------------------------

axes[0,0].imshow(
    image,
    cmap='gray'
)

axes[0,0].set_title(
    "Input Image"
)

axes[0,0].axis('off')

# ------------------------------------------
# Feature Map 3x3
# ------------------------------------------

f3 = feature3.squeeze().numpy()

im1 = axes[0,1].imshow(
    f3,
    cmap='jet'
)

axes[0,1].set_title(
    "Feature Map (3x3 Kernel)"
)

axes[0,1].axis('off')

for y in range(f3.shape[0]):
    for x in range(f3.shape[1]):
        axes[0,1].text(
            x,
            y,
            int(f3[y,x]),
            ha='center',
            va='center',
            color='white',
            fontsize=8
        )

# ------------------------------------------
# Feature Map 5x5
# ------------------------------------------

f5 = feature5.squeeze().numpy()

im2 = axes[1,0].imshow(
    f5,
    cmap='jet'
)

axes[1,0].set_title(
    "Feature Map (5x5 Kernel)"
)

axes[1,0].axis('off')

for y in range(f5.shape[0]):
    for x in range(f5.shape[1]):
        axes[1,0].text(
            x,
            y,
            int(f5[y,x]),
            ha='center',
            va='center',
            color='white',
            fontsize=8
        )

# ------------------------------------------
# Feature Map 7x7
# ------------------------------------------

f7 = feature7.squeeze().numpy()

im3 = axes[1,1].imshow(
    f7,
    cmap='jet'
)

axes[1,1].set_title(
    "Feature Map (7x7 Kernel)"
)

axes[1,1].axis('off')

for y in range(f7.shape[0]):
    for x in range(f7.shape[1]):
        axes[1,1].text(
            x,
            y,
            int(f7[y,x]),
            ha='center',
            va='center',
            color='white',
            fontsize=8
        )

plt.tight_layout()
plt.show()