# 卷積運算：用公式手算，和 PyTorch 的 conv2d 比對；並檢查輸出大小公式 (n + 2p − k)/s + 1
import torch
import torch.nn.functional as F

x = torch.tensor([[1., 2., 0., 1., 3.],
                  [0., 1., 3., 2., 1.],
                  [2., 0., 1., 1., 0.],
                  [1., 3., 2., 0., 2.],
                  [0., 1., 1., 2., 1.]])
w = torch.tensor([[1., 0., -1.],
                  [2., 0., -2.],
                  [1., 0., -1.]])                # Sobel 核：偵測水平方向的亮度變化（垂直邊緣）

# 公式：y[i, j] = Σ_m Σ_n w[m, n] · x[i + m, j + n]（深度學習框架實際算的是互相關，不翻轉卷積核）
y = torch.zeros(3, 3)
for i in range(3):
    for j in range(3):
        y[i, j] = (w * x[i:i + 3, j:j + 3]).sum()
print("手算：\n", y)

y2 = F.conv2d(x.view(1, 1, 5, 5), w.view(1, 1, 3, 3)).squeeze()   # 形狀＝（批次, 通道, 高, 寬）
print("conv2d 與手算相同：", torch.equal(y, y2))

for n, k, p, s in [(28, 5, 0, 1), (28, 3, 1, 1), (32, 3, 1, 2), (224, 7, 3, 2)]:
    out = F.conv2d(torch.zeros(1, 1, n, n), torch.zeros(1, 1, k, k), padding=p, stride=s).shape[-1]
    print(f"n={n:3d} k={k} p={p} s={s} → 輸出 {out:3d}，公式 {(n + 2 * p - k) // s + 1:3d}")