卷积神经网络 CNN

"""图像在内存里长什么样：像素、通道、HWC。"""
import numpy as np
import matplotlib.pyplot as plt

# 像素值 0 最暗（黑），255 最亮（白）
black = np.zeros([200, 200, 3])            # 全 0 -> 纯黑
white = np.full([200, 200, 3], 255)        # 全 255 -> 纯白
plt.imshow(black); plt.show()
plt.imshow(white); plt.show()

# 读一张真实图片：matplotlib 给出的顺序是 (高, 宽, 通道)，简称 HWC
img = plt.imread('data/img.jpg')
print('图像形状 (H, W, C):', img.shape)     # 例如 (640, 640, 3)
print('取值范围:', img.min(), '~', img.max())

plt.imshow(img)
plt.axis('off')
plt.show()

# PyTorch 的卷积层要的是 (N, C, H, W)，和这里差一次转置，别记混

"""全连接层处理图像为什么不可行：参数量对着算一遍就知道了。

不依赖任何第三方库，纯算术，可以直接跑出确定的数字。

两个公式：
    全连接层参数量 = 输入元素个数 × 输出神经元个数 + 输出神经元个数(偏置)
    卷积层参数量   = 卷积核个数 × (输入通道数 × 核高 × 核宽 + 1)

关键差别在于：全连接的参数量**随图片尺寸平方级增长**，
卷积的参数量**和图片多大完全无关**——同一枚印章走遍全图。
"""


def fc_params(in_elements, out_units):
    """一个全连接层的参数量。"""
    return in_elements * out_units + out_units


def conv_params(in_channels, out_channels, ksize):
    """一个卷积层的参数量。注意式子里没有 H 和 W。"""
    return out_channels * (in_channels * ksize * ksize + 1)


def human(n):
    """把参数量换算成更好理解的量级。"""
    if n >= 1e8:
        return '%.2f 亿' % (n / 1e8)
    if n >= 1e4:
        return '%.2f 万' % (n / 1e4)
    return '%d' % n


def compare(size, channels=3, hidden=1000, out_channels=6, ksize=3):
    in_elements = channels * size * size
    fc = fc_params(in_elements, hidden)
    conv = conv_params(channels, out_channels, ksize)
    return in_elements, fc, conv


if __name__ == '__main__':
    print('第一个隐藏层取 1000 个神经元，卷积层取 6 个 3×3 卷积核：\n')
    print('%-12s %-14s %-18s %-14s %s'
          % ('图片边长', '拉平后长度', '全连接参数量', '卷积参数量', '倍数'))
    for size in (32, 64, 224, 512, 1024):
        in_elements, fc, conv = compare(size)
        print('%-12s %-14s %-18s %-14s %.0f 倍'
              % ('%d×%d' % (size, size), human(in_elements),
                 human(fc), human(conv), fc / conv))

    print('\n看最后一列：图片边长翻一倍，全连接参数量翻四倍，卷积那一列纹丝不动。')

    # ------------------------------------------------------------ 第二个问题
    print('\n全连接的第二宗罪：把二维结构拍扁之后，「相邻」这件事就丢了。')
    img = [[1, 2, 3],
           [4, 5, 6],
           [7, 8, 9]]
    flat = [v for row in img for v in row]
    print('原图 3×3 ：', img)
    print('拉平之后 ：', flat)
    print('像素 2 在原图里上下左右分别是 %s；拉平后它的邻居只剩下 %s。'
          % ([img[1][1]], [flat[0], flat[2]]))
    print('也就是说：5 明明就在 2 的正下方，拉平后却隔了 2 个位置，')
    print('全连接层无从知道这层关系，只能从数据里硬学——这正是卷积要替它做的事。')

    # ------------------------------------------------------------ 真实网络对账
    print('\n本页 CIFAR-10 网络的卷积层参数量（和图片尺寸无关的那部分）：')
    c1 = conv_params(3, 6, 3)
    c2 = conv_params(6, 16, 3)
    print('  卷积1  3 -> 6  通道，3×3 核：6 × (3×3×3 + 1) = %d' % c1)
    print('  卷积2  6 -> 16 通道，3×3 核：16 × (6×3×3 + 1) = %d' % c2)
    print('  两个卷积层合计 %d 个参数' % (c1 + c2))
    fc1 = fc_params(576, 120)
    print('  而紧随其后的一个全连接层 576 -> 120 就要 %d 个参数' % fc1)
    print('  比例：%.1f 倍。CNN 的参数大头一直在全连接层，不在卷积层。'
          % (fc1 / (c1 + c2)))

"""手写卷积：不用 torch，把「对应相乘再相加」逐格算出来。

跑一遍这份代码，卷积就不再是黑箱了——它只是一个双重循环，
每一步取出一个和卷积核同样大小的窗口，对应位置相乘、全部加起来。

顺带把输出尺寸公式 N = (W - F + 2P) / S + 1 用实际循环次数验证一遍：
循环真正跑了多少行多少列，公式就必须给出同样的数。
"""


def conv2d(image, kernel, stride=1, padding=0, bias=0.0):
    """单通道二维卷积。image / kernel 都是二维嵌套列表。"""
    if padding > 0:
        image = pad_zeros(image, padding)

    in_h, in_w = len(image), len(image[0])
    k_h, k_w = len(kernel), len(kernel[0])

    out = []
    row = 0
    while row + k_h <= in_h:                 # 纵向滑动
        line = []
        col = 0
        while col + k_w <= in_w:             # 横向滑动
            total = 0.0
            for i in range(k_h):             # 窗口内逐格相乘再相加
                for j in range(k_w):
                    total += image[row + i][col + j] * kernel[i][j]
            line.append(total + bias)        # 每个卷积核配一个偏置
            col += stride
        out.append(line)
        row += stride
    return out


def pad_zeros(image, p):
    """四周补 p 圈 0。补完边长增加 2p，这就是公式里 +2P 的来历。"""
    width = len(image[0]) + 2 * p
    top = [[0] * width for _ in range(p)]
    middle = [[0] * p + list(r) + [0] * p for r in image]
    bottom = [[0] * width for _ in range(p)]
    return top + middle + bottom


def out_size(w, f, p, s):
    """输出尺寸公式；除不尽时向下取整（窗口伸出边界就不算这一步）。"""
    return (w - f + 2 * p) // s + 1


def show(name, mat):
    print('%s  (%d×%d)' % (name, len(mat), len(mat[0])))
    for row in mat:
        print('   ' + ' '.join('%6.1f' % v for v in row))


if __name__ == '__main__':
    # 6×6 的输入，格子里填 1..36 方便肉眼核对
    image = [[r * 6 + c + 1 for c in range(6)] for r in range(6)]
    show('输入图像', image)

    # 一枚 3×3 的「印章」：中间重、四周轻，作用是突出中心像素
    kernel = [[0, -1, 0],
              [-1, 5, -1],
              [0, -1, 0]]
    show('\n卷积核', kernel)

    # ---- 步长 1、不补零：6×6 -> 4×4，正是图①画的那个情形
    feat = conv2d(image, kernel, stride=1, padding=0)
    show('\n步长1 不补零 的特征图', feat)
    print('   公式预测边长 = (6 - 3 + 2×0) / 1 + 1 = %d，实测 %d ✔'
          % (out_size(6, 3, 0, 1), len(feat)))

    # 手算第一格验证：窗口是 [[1,2,3],[7,8,9],[13,14,15]]
    manual = (1 * 0 + 2 * (-1) + 3 * 0
              + 7 * (-1) + 8 * 5 + 9 * (-1)
              + 13 * 0 + 14 * (-1) + 15 * 0)
    print('   左上角那一格手算 = %d，代码算出 = %.0f' % (manual, feat[0][0]))

    # ---- 步长 2：窗口一次跨两格，输出边长直接减半
    feat2 = conv2d(image, kernel, stride=2, padding=0)
    show('\n步长2 不补零 的特征图', feat2)
    print('   公式预测 = (6 - 3 + 0) / 2 + 1 = %d，实测 %d ✔'
          % (out_size(6, 3, 0, 2), len(feat2)))

    # ---- 补一圈 0：尺寸被保住
    small = [[r * 5 + c + 1 for c in range(5)] for r in range(5)]
    feat3 = conv2d(small, kernel, stride=1, padding=1)
    print('\n5×5 输入，补一圈 0，步长 1 -> 输出 %d×%d' % (len(feat3), len(feat3[0])))
    print('   公式预测 = (5 - 3 + 2×1) / 1 + 1 = %d ✔' % out_size(5, 3, 1, 1))
    print('   这就是「四周补一圈 0 尺寸就保住」的由来：3×3 核配 padding=1。')

    # ---- 参数量：和图片多大没有关系
    print('\n这枚 3×3 卷积核的参数量 = 3×3 + 1 = 10 个，')
    print('无论输入是 6×6 还是 1024×1024，都是这 10 个数在反复使用。')

"""padding 与 stride 的组合实验：改一个参数，看特征图尺寸怎么变。

每组都先写出公式预测值，再让 PyTorch 实际算一遍，两边必须一致。
公式：N = (W - F + 2P) / S + 1，除不尽向下取整。
"""
import torch
import torch.nn as nn


def predict(w, f, p, s):
    """公式预测的输出边长。"""
    return (w - f + 2 * p) // s + 1


def check(in_size, ksize, padding, stride, in_ch=3, out_ch=8):
    """建一个卷积层实跑一次，和公式对账。"""
    conv = nn.Conv2d(in_ch, out_ch, kernel_size=ksize,
                     stride=stride, padding=padding)
    x = torch.randn(1, in_ch, in_size, in_size)
    out = conv(x)

    expect = predict(in_size, ksize, padding, stride)
    actual = out.shape[-1]
    flag = '一致' if expect == actual else '不一致'
    print('W=%-4d F=%d P=%d S=%d  公式 %-4d 实测 %-4d  %s  输出形状 %s'
          % (in_size, ksize, padding, stride, expect, actual, flag, tuple(out.shape)))
    return actual


if __name__ == '__main__':
    print('一、padding 的作用：把被卷掉的边缘补回来')
    check(5, 3, padding=0, stride=1)      # 5 -> 3，越卷越小
    check(5, 3, padding=1, stride=1)      # 5 -> 5，尺寸保住
    check(32, 3, padding=0, stride=1)     # 32 -> 30
    check(32, 3, padding=1, stride=1)     # 32 -> 32

    print('\n  规律：核边长 F 为奇数时，padding = (F-1)/2 正好让尺寸不变。')
    for f in (1, 3, 5, 7):
        p = (f - 1) // 2
        print('    F=%d -> padding=%d，此时 32×32 卷完还是 %d×%d'
              % (f, p, predict(32, f, p, 1), predict(32, f, p, 1)))

    print('\n二、stride 的作用：一步跨多远，直接决定降维速度')
    for s in (1, 2, 3):
        check(32, 3, padding=1, stride=s)
    print('  步长每翻一倍，输出边长大致减半，计算量减到四分之一。')

    print('\n三、除不尽的情况：向下取整，右下角那一条被丢掉')
    check(7, 3, padding=0, stride=2)      # (7-3)/2+1 = 3
    check(8, 3, padding=0, stride=2)      # (8-3)//2+1 = 3，和 7 一样
    print('  注意上面两行输出边长相同：边长 8 的那一列多出来的像素没人管，')
    print('  想一个不漏就得补 padding，或者让尺寸能被整除。')

    print('\n四、池化层用的是同一个公式')
    pool = nn.MaxPool2d(kernel_size=2, stride=2)
    x = torch.randn(1, 6, 30, 30)
    print('  6×30×30 经过 2×2 步长 2 的最大池化 ->', tuple(pool(x).shape))
    print('  公式：(30 - 2 + 0) / 2 + 1 = %d ' % predict(30, 2, 0, 2))
    print('  通道数 6 没变——池化对每个通道单独做。')

    print('\n五、参数量只和核大小、通道数有关，和输入尺寸无关')
    for size in (32, 224):
        conv = nn.Conv2d(3, 6, kernel_size=3)
        n = sum(p.numel() for p in conv.parameters())
        print('  输入 %d×%d 时，这个卷积层参数量 = %d' % (size, size, n))
    print('  6 × (3×3×3 + 1) = 168，两次都一样。')

"""逐层推导特征图尺寸：把「576 是怎么来的」算给自己看。

不依赖 torch，纯算术。搭网络时先用它把每一层的形状过一遍，
比搭完跑起来再看报错快得多。

核心公式（卷积和池化共用同一个）：
    N = (W - F + 2P) / S + 1     除不尽时向下取整
其中 W 输入边长、F 核边长、P padding 圈数、S 步长。
"""


def out_size(w, f, p=0, s=1):
    """算一层之后的边长。向下取整：窗口伸出边界的那一步不算数。"""
    return (w - f + 2 * p) // s + 1


class Shape:
    """记录 (通道数, 高, 宽) 并逐层推进。"""

    def __init__(self, c, h, w, name='输入'):
        self.c, self.h, self.w = c, h, w
        self.log = ['%-10s %d×%d×%d' % (name, c, h, w)]

    def conv(self, out_channels, ksize, stride=1, padding=0, name='卷积'):
        self.h = out_size(self.h, ksize, padding, stride)
        self.w = out_size(self.w, ksize, padding, stride)
        self.c = out_channels                    # 卷积核有几个，输出就有几个通道
        self.log.append('%-10s %d×%d×%d   （核 %d×%d，步长 %d，padding %d）'
                        % (name, self.c, self.h, self.w, ksize, ksize, stride, padding))
        return self

    def pool(self, ksize=2, stride=2, padding=0, name='池化'):
        self.h = out_size(self.h, ksize, padding, stride)
        self.w = out_size(self.w, ksize, padding, stride)
        # 通道数不变：池化对每个通道单独做，不像卷积会把通道加起来
        self.log.append('%-10s %d×%d×%d   （窗口 %d×%d，步长 %d）'
                        % (name, self.c, self.h, self.w, ksize, ksize, stride))
        return self

    def flatten(self):
        n = self.c * self.h * self.w
        self.log.append('%-10s %d   （= %d × %d × %d）'
                        % ('拉平', n, self.c, self.h, self.w))
        return n

    def show(self):
        for line in self.log:
            print('  ' + line)


def conv_params(in_c, out_c, k):
    return out_c * (in_c * k * k + 1)


def fc_params(in_f, out_f):
    return in_f * out_f + out_f


if __name__ == '__main__':
    print('一、公式自检（几组能手算的例子）')
    cases = [
        (5, 3, 1, 1, '补一圈 0，尺寸保住'),
        (5, 3, 0, 1, '不补零，越卷越小'),
        (6, 3, 0, 1, '印章图里的 6×6 -> 4×4'),
        (32, 3, 0, 1, 'CIFAR 第一层卷积'),
        (30, 2, 0, 2, 'CIFAR 第一层池化'),
        (7, 3, 0, 2, '步长 2，除不尽向下取整'),
    ]
    for w, f, p, s, note in cases:
        print('  W=%-4d F=%d P=%d S=%d -> N=%-4d %s'
              % (w, f, p, s, out_size(w, f, p, s), note))

    print('\n二、逐层推导 CIFAR-10 分类网络')
    shape = Shape(3, 32, 32)
    shape.conv(6, ksize=3, stride=1, padding=0, name='卷积1')
    shape.pool(2, 2, name='池化1')
    shape.conv(16, ksize=3, stride=1, padding=0, name='卷积2')
    shape.pool(2, 2, name='池化2')
    flat = shape.flatten()
    shape.show()

    print('\n  拉平后维度 = %d —— nn.Linear 的第一个参数必须写这个数' % flat)
    assert flat == 576, '和网络里写的 576 对不上，说明哪一层参数改了'
    print('  断言通过：与网络里写死的 576 一致。')

    print('\n三、这个网络的参数量分布')
    p_conv1 = conv_params(3, 6, 3)
    p_conv2 = conv_params(6, 16, 3)
    p_fc1 = fc_params(576, 120)
    p_fc2 = fc_params(120, 84)
    p_out = fc_params(84, 10)
    total = p_conv1 + p_conv2 + p_fc1 + p_fc2 + p_out
    for name, p in (('卷积1', p_conv1), ('卷积2', p_conv2), ('全连接 576→120', p_fc1),
                    ('全连接 120→84', p_fc2), ('输出 84→10', p_out)):
        print('  %-16s %7d 个  占比 %5.1f%%' % (name, p, 100.0 * p / total))
    print('  合计 %d 个参数' % total)
    print('  池化层 0 个参数——它只挑数不学东西。')

    print('\n四、改一个参数会连累后面所有层')
    print('  把两个卷积核都换成 5×5，其余不动：')
    alt = Shape(3, 32, 32)
    alt.conv(6, ksize=5, name='卷积1')
    alt.pool(2, 2, name='池化1')
    alt.conv(16, ksize=5, name='卷积2')
    alt.pool(2, 2, name='池化2')
    alt_flat = alt.flatten()
    alt.show()
    print('  拉平后变成 %d，全连接层第一个参数就得跟着改成 %d。' % (alt_flat, alt_flat))
    print('  忘了改的现象：mat1 and mat2 shapes cannot be multiplied。')

"""多通道与多卷积核：通道数是怎么一步步变过去的。

两条规则，记住就不会算错：
  1. 一个卷积核的**深度必须等于输入通道数**——它要同时盖住所有通道，
     各通道分别相乘求和，最后加到一起，所以**一个核只产出一张特征图**。
  2. 用几个卷积核，输出就有几个通道。out_channels 就是卷积核的个数。

于是通道数的变化完全由 out_channels 决定，与输入通道数无关。
"""
import torch
import torch.nn as nn


def show(name, tensor):
    print('%-28s %s' % (name, tuple(tensor.shape)))


if __name__ == '__main__':
    torch.manual_seed(0)

    # ---------------------------------------------------------- 单通道 -> 单核
    print('一、最简单的情形：1 个通道，1 个卷积核')
    x1 = torch.randn(1, 1, 6, 6)          # (N, C, H, W)
    conv1 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=3)
    show('输入', x1)
    show('输出（一张特征图）', conv1(x1))
    print('  权重形状', tuple(conv1.weight.shape), '= (核个数, 输入通道, 核高, 核宽)')
    print('  参数量 = 1 × (1×3×3 + 1) =', sum(p.numel() for p in conv1.parameters()))

    # ---------------------------------------------------------- 多通道 -> 单核
    print('\n二、彩色图 3 通道，仍然只用 1 个卷积核')
    x3 = torch.randn(1, 3, 6, 6)
    conv2 = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3)
    show('输入', x3)
    show('输出（还是一张特征图）', conv2(x3))
    print('  权重形状', tuple(conv2.weight.shape), '—— 核的深度自动变成 3')
    print('  三个通道各自卷完之后**相加**成一张，所以输出通道数还是 1。')
    print('  参数量 = 1 × (3×3×3 + 1) =', sum(p.numel() for p in conv2.parameters()))

    # ---------------------------------------------------------- 多通道 -> 多核
    print('\n三、3 通道输入，6 个卷积核')
    conv3 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3)
    show('输入', x3)
    show('输出（6 张特征图摞起来）', conv3(x3))
    print('  权重形状', tuple(conv3.weight.shape))
    print('  参数量 = 6 × (3×3×3 + 1) =', sum(p.numel() for p in conv3.parameters()))
    print('  每个核学一种花纹：有的对横边敏感，有的对竖边敏感，有的对颜色块敏感。')

    # ---------------------------------------------------------- 通道数错配
    print('\n四、in_channels 写错会怎样')
    wrong = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=3)
    try:
        wrong(x3)                          # 输入 3 通道，层却声明 1 通道
    except RuntimeError as exc:
        print('  报错：', str(exc).split('\n')[0])
    print('  修法：in_channels 必须等于上一层的输出通道数，第一层等于图片通道数。')

    # ---------------------------------------------------------- 池化不改通道
    print('\n五、对比：池化层不会合并通道')
    feat = conv3(x3)
    pooled = nn.MaxPool2d(kernel_size=2, stride=2)(feat)
    show('卷积输出', feat)
    show('池化之后', pooled)
    print('  通道数 6 原封不动，只有高和宽减半——这是池化和卷积最本质的区别。')

    # ---------------------------------------------------------- 整条链路
    print('\n六、把两层串起来，看通道怎么一路变过去')
    net = nn.Sequential(
        nn.Conv2d(3, 6, kernel_size=3), nn.ReLU(),
        nn.MaxPool2d(2, 2),
        nn.Conv2d(6, 16, kernel_size=3), nn.ReLU(),
        nn.MaxPool2d(2, 2),
    )
    out = net(torch.randn(1, 3, 32, 32))
    show('3×32×32 走完两组卷积池化', out)
    print('  拉平长度 =', out.numel(), '= 16 × 6 × 6')

"""池化层：只降维，不带参数，逐通道各池各的。"""
import torch
import torch.nn as nn

# ---------- 单通道 ----------
x = torch.tensor([[[0, 1, 2],
                   [3, 4, 5],
                   [6, 7, 8]]]).float()      # (1, 3, 3)

print('最大池化:\n', nn.MaxPool2d(kernel_size=2, stride=1, padding=0)(x))
# 每个 2×2 窗口取最大：[[4,5],[7,8]]

print('平均池化:\n', nn.AvgPool2d(kernel_size=2, stride=1, padding=0)(x))
# 每个 2×2 窗口取均值：[[2,3],[5,6]]

# ---------- 多通道 ----------
x3 = torch.tensor([[[0, 1, 2], [3, 4, 5], [6, 7, 8]],
                   [[10, 20, 30], [40, 50, 60], [70, 80, 90]],
                   [[11, 22, 33], [44, 55, 66], [77, 88, 99]]]).float()
out = nn.MaxPool2d(kernel_size=2, stride=1, padding=0)(x3)
print('多通道池化后形状', out.shape)          # (3, 2, 2)
# 通道数不变：池化对每个通道单独做，不像卷积会把通道加起来

# 池化层没有任何可学习参数
print('池化层参数量', sum(p.numel() for p in nn.MaxPool2d(2).parameters()))   # 0

"""卷积层最小例子：一张图进去，一叠特征图出来。"""
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

img = plt.imread('data/img.jpg')            # (640, 640, 3)，HWC

# HWC -> CHW，再在最前面补一个 batch 维，变成 (1, 3, 640, 640)
x = torch.tensor(img).permute(2, 0, 1).unsqueeze(0).to(torch.float32)
print('送进卷积层的形状', x.shape)

# in_channels 必须等于输入的通道数；out_channels 就是卷积核的个数
conv = nn.Conv2d(in_channels=3, out_channels=3,
                 kernel_size=3, stride=2, padding=0)

feature_map = conv(x)
print('特征图形状', feature_map.shape)
# 手算验证：N = (640 - 3 + 2*0) / 2 + 1 = 319（除不尽时向下取整）

# 卷积层的参数量 = 卷积核个数 × (通道数 × 核高 × 核宽 + 1 个偏置)
#               = 3 × (3 × 3 × 3 + 1) = 84，和图片多大完全无关
print('参数量', sum(p.numel() for p in conv.parameters()))

"""池化层：只降维，不带参数，逐通道各池各的。"""
import torch
import torch.nn as nn

# ---------- 单通道 ----------
x = torch.tensor([[[0, 1, 2],
                   [3, 4, 5],
                   [6, 7, 8]]]).float()      # (1, 3, 3)

print('最大池化:\n', nn.MaxPool2d(kernel_size=2, stride=1, padding=0)(x))
# 每个 2×2 窗口取最大：[[4,5],[7,8]]

print('平均池化:\n', nn.AvgPool2d(kernel_size=2, stride=1, padding=0)(x))
# 每个 2×2 窗口取均值：[[2,3],[5,6]]

# ---------- 多通道 ----------
x3 = torch.tensor([[[0, 1, 2], [3, 4, 5], [6, 7, 8]],
                   [[10, 20, 30], [40, 50, 60], [70, 80, 90]],
                   [[11, 22, 33], [44, 55, 66], [77, 88, 99]]]).float()
out = nn.MaxPool2d(kernel_size=2, stride=1, padding=0)(x3)
print('多通道池化后形状', out.shape)          # (3, 2, 2)
# 通道数不变：池化对每个通道单独做，不像卷积会把通道加起来

# 池化层没有任何可学习参数
print('池化层参数量', sum(p.numel() for p in nn.MaxPool2d(2).parameters()))   # 0

"""手写卷积：不用 torch，把「对应相乘再相加」逐格算出来。

跑一遍这份代码，卷积就不再是黑箱了——它只是一个双重循环，
每一步取出一个和卷积核同样大小的窗口，对应位置相乘、全部加起来。

顺带把输出尺寸公式 N = (W - F + 2P) / S + 1 用实际循环次数验证一遍：
循环真正跑了多少行多少列，公式就必须给出同样的数。
"""


def conv2d(image, kernel, stride=1, padding=0, bias=0.0):
    """单通道二维卷积。image / kernel 都是二维嵌套列表。"""
    if padding > 0:
        image = pad_zeros(image, padding)

    in_h, in_w = len(image), len(image[0])
    k_h, k_w = len(kernel), len(kernel[0])

    out = []
    row = 0
    while row + k_h <= in_h:                 # 纵向滑动
        line = []
        col = 0
        while col + k_w <= in_w:             # 横向滑动
            total = 0.0
            for i in range(k_h):             # 窗口内逐格相乘再相加
                for j in range(k_w):
                    total += image[row + i][col + j] * kernel[i][j]
            line.append(total + bias)        # 每个卷积核配一个偏置
            col += stride
        out.append(line)
        row += stride
    return out


def pad_zeros(image, p):
    """四周补 p 圈 0。补完边长增加 2p，这就是公式里 +2P 的来历。"""
    width = len(image[0]) + 2 * p
    top = [[0] * width for _ in range(p)]
    middle = [[0] * p + list(r) + [0] * p for r in image]
    bottom = [[0] * width for _ in range(p)]
    return top + middle + bottom


def out_size(w, f, p, s):
    """输出尺寸公式；除不尽时向下取整（窗口伸出边界就不算这一步）。"""
    return (w - f + 2 * p) // s + 1


def show(name, mat):
    print('%s  (%d×%d)' % (name, len(mat), len(mat[0])))
    for row in mat:
        print('   ' + ' '.join('%6.1f' % v for v in row))


if __name__ == '__main__':
    # 6×6 的输入，格子里填 1..36 方便肉眼核对
    image = [[r * 6 + c + 1 for c in range(6)] for r in range(6)]
    show('输入图像', image)

    # 一枚 3×3 的「印章」：中间重、四周轻，作用是突出中心像素
    kernel = [[0, -1, 0],
              [-1, 5, -1],
              [0, -1, 0]]
    show('\n卷积核', kernel)

    # ---- 步长 1、不补零：6×6 -> 4×4，正是图①画的那个情形
    feat = conv2d(image, kernel, stride=1, padding=0)
    show('\n步长1 不补零 的特征图', feat)
    print('   公式预测边长 = (6 - 3 + 2×0) / 1 + 1 = %d，实测 %d ✔'
          % (out_size(6, 3, 0, 1), len(feat)))

    # 手算第一格验证：窗口是 [[1,2,3],[7,8,9],[13,14,15]]
    manual = (1 * 0 + 2 * (-1) + 3 * 0
              + 7 * (-1) + 8 * 5 + 9 * (-1)
              + 13 * 0 + 14 * (-1) + 15 * 0)
    print('   左上角那一格手算 = %d，代码算出 = %.0f' % (manual, feat[0][0]))

    # ---- 步长 2：窗口一次跨两格，输出边长直接减半
    feat2 = conv2d(image, kernel, stride=2, padding=0)
    show('\n步长2 不补零 的特征图', feat2)
    print('   公式预测 = (6 - 3 + 0) / 2 + 1 = %d，实测 %d ✔'
          % (out_size(6, 3, 0, 2), len(feat2)))

    # ---- 补一圈 0：尺寸被保住
    small = [[r * 5 + c + 1 for c in range(5)] for r in range(5)]
    feat3 = conv2d(small, kernel, stride=1, padding=1)
    print('\n5×5 输入，补一圈 0，步长 1 -> 输出 %d×%d' % (len(feat3), len(feat3[0])))
    print('   公式预测 = (5 - 3 + 2×1) / 1 + 1 = %d ✔' % out_size(5, 3, 1, 1))
    print('   这就是「四周补一圈 0 尺寸就保住」的由来：3×3 核配 padding=1。')

    # ---- 参数量：和图片多大没有关系
    print('\n这枚 3×3 卷积核的参数量 = 3×3 + 1 = 10 个，')
    print('无论输入是 6×6 还是 1024×1024，都是这 10 个数在反复使用。')

"""逐层推导特征图尺寸：把「576 是怎么来的」算给自己看。

不依赖 torch，纯算术。搭网络时先用它把每一层的形状过一遍，
比搭完跑起来再看报错快得多。

核心公式（卷积和池化共用同一个）：
    N = (W - F + 2P) / S + 1     除不尽时向下取整
其中 W 输入边长、F 核边长、P padding 圈数、S 步长。
"""


def out_size(w, f, p=0, s=1):
    """算一层之后的边长。向下取整：窗口伸出边界的那一步不算数。"""
    return (w - f + 2 * p) // s + 1


class Shape:
    """记录 (通道数, 高, 宽) 并逐层推进。"""

    def __init__(self, c, h, w, name='输入'):
        self.c, self.h, self.w = c, h, w
        self.log = ['%-10s %d×%d×%d' % (name, c, h, w)]

    def conv(self, out_channels, ksize, stride=1, padding=0, name='卷积'):
        self.h = out_size(self.h, ksize, padding, stride)
        self.w = out_size(self.w, ksize, padding, stride)
        self.c = out_channels                    # 卷积核有几个，输出就有几个通道
        self.log.append('%-10s %d×%d×%d   （核 %d×%d，步长 %d，padding %d）'
                        % (name, self.c, self.h, self.w, ksize, ksize, stride, padding))
        return self

    def pool(self, ksize=2, stride=2, padding=0, name='池化'):
        self.h = out_size(self.h, ksize, padding, stride)
        self.w = out_size(self.w, ksize, padding, stride)
        # 通道数不变：池化对每个通道单独做，不像卷积会把通道加起来
        self.log.append('%-10s %d×%d×%d   （窗口 %d×%d，步长 %d）'
                        % (name, self.c, self.h, self.w, ksize, ksize, stride))
        return self

    def flatten(self):
        n = self.c * self.h * self.w
        self.log.append('%-10s %d   （= %d × %d × %d）'
                        % ('拉平', n, self.c, self.h, self.w))
        return n

    def show(self):
        for line in self.log:
            print('  ' + line)


def conv_params(in_c, out_c, k):
    return out_c * (in_c * k * k + 1)


def fc_params(in_f, out_f):
    return in_f * out_f + out_f


if __name__ == '__main__':
    print('一、公式自检（几组能手算的例子）')
    cases = [
        (5, 3, 1, 1, '补一圈 0，尺寸保住'),
        (5, 3, 0, 1, '不补零，越卷越小'),
        (6, 3, 0, 1, '印章图里的 6×6 -> 4×4'),
        (32, 3, 0, 1, 'CIFAR 第一层卷积'),
        (30, 2, 0, 2, 'CIFAR 第一层池化'),
        (7, 3, 0, 2, '步长 2，除不尽向下取整'),
    ]
    for w, f, p, s, note in cases:
        print('  W=%-4d F=%d P=%d S=%d -> N=%-4d %s'
              % (w, f, p, s, out_size(w, f, p, s), note))

    print('\n二、逐层推导 CIFAR-10 分类网络')
    shape = Shape(3, 32, 32)
    shape.conv(6, ksize=3, stride=1, padding=0, name='卷积1')
    shape.pool(2, 2, name='池化1')
    shape.conv(16, ksize=3, stride=1, padding=0, name='卷积2')
    shape.pool(2, 2, name='池化2')
    flat = shape.flatten()
    shape.show()

    print('\n  拉平后维度 = %d —— nn.Linear 的第一个参数必须写这个数' % flat)
    assert flat == 576, '和网络里写的 576 对不上，说明哪一层参数改了'
    print('  断言通过：与网络里写死的 576 一致。')

    print('\n三、这个网络的参数量分布')
    p_conv1 = conv_params(3, 6, 3)
    p_conv2 = conv_params(6, 16, 3)
    p_fc1 = fc_params(576, 120)
    p_fc2 = fc_params(120, 84)
    p_out = fc_params(84, 10)
    total = p_conv1 + p_conv2 + p_fc1 + p_fc2 + p_out
    for name, p in (('卷积1', p_conv1), ('卷积2', p_conv2), ('全连接 576→120', p_fc1),
                    ('全连接 120→84', p_fc2), ('输出 84→10', p_out)):
        print('  %-16s %7d 个  占比 %5.1f%%' % (name, p, 100.0 * p / total))
    print('  合计 %d 个参数' % total)
    print('  池化层 0 个参数——它只挑数不学东西。')

    print('\n四、改一个参数会连累后面所有层')
    print('  把两个卷积核都换成 5×5，其余不动：')
    alt = Shape(3, 32, 32)
    alt.conv(6, ksize=5, name='卷积1')
    alt.pool(2, 2, name='池化1')
    alt.conv(16, ksize=5, name='卷积2')
    alt.pool(2, 2, name='池化2')
    alt_flat = alt.flatten()
    alt.show()
    print('  拉平后变成 %d，全连接层第一个参数就得跟着改成 %d。' % (alt_flat, alt_flat))
    print('  忘了改的现象：mat1 and mat2 shapes cannot be multiplied。')

"""完整案例：CIFAR10 十分类。

5 万张 32×32 彩色训练图、1 万张测试图、10 个类别。
网络：卷积 -> 池化 -> 卷积 -> 池化 -> 拉平 -> 三个全连接。
"""
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor

BATCH_SIZE = 8


def create_dataset(root='data'):
    # ToTensor 做两件事：HWC -> CHW，并把像素从 0~255 归一到 0.0~1.0
    train = CIFAR10(root=root, train=True, transform=ToTensor(), download=True)
    valid = CIFAR10(root=root, train=False, transform=ToTensor(), download=True)
    return train, valid


class ImageClassification(nn.Module):
    def __init__(self):
        super().__init__()
        # 输入 3×32×32
        self.conv1 = nn.Conv2d(3, 6, kernel_size=3, stride=1)    # -> 6×30×30
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)       # -> 6×15×15
        self.conv2 = nn.Conv2d(6, 16, kernel_size=3, stride=1)   # -> 16×13×13
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)       # -> 16×6×6

        # 576 = 16 × 6 × 6，这个数是上面一路算下来的，不能拍脑袋写
        self.linear1 = nn.Linear(576, 120)
        self.linear2 = nn.Linear(120, 84)
        self.out = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv1(x)))
        x = self.pool2(torch.relu(self.conv2(x)))
        # 把 (batch, 16, 6, 6) 摊平成 (batch, 576) 才能进全连接层
        x = x.reshape(x.size(0), -1)
        x = torch.relu(self.linear1(x))
        x = torch.relu(self.linear2(x))
        return self.out(x)          # 原始 logits，softmax 交给损失函数


def train(model, train_dataset, epochs=10, lr=1e-3):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    for epoch in range(epochs):
        dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
        total_loss, sample_num, start = 0.0, 0, time.time()

        model.train()
        for x, y in dataloader:
            loss = criterion(model(x), y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * len(y)
            sample_num += len(y)

        print('epoch %2d  loss %.5f  time %.2fs'
              % (epoch + 1, total_loss / sample_num, time.time() - start))

    torch.save(model.state_dict(), 'model/image_classification.pth')
    return model


def evaluate(valid_dataset, weight='model/image_classification.pth'):
    model = ImageClassification()
    model.load_state_dict(torch.load(weight))
    model.eval()

    dataloader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False)
    correct, total = 0, 0
    with torch.no_grad():
        for x, y in dataloader:
            correct += (model(x).argmax(dim=-1) == y).sum().item()
            total += len(y)
    print('测试集准确率 %.4f' % (correct / total))


def main():
    train_dataset, valid_dataset = create_dataset()
    print('训练集', train_dataset.data.shape)     # (50000, 32, 32, 3)
    print('测试集', valid_dataset.data.shape)     # (10000, 32, 32, 3)
    print('类别表', train_dataset.class_to_idx)

    model = ImageClassification()
    train(model, train_dataset)
    evaluate(valid_dataset)


if __name__ == '__main__':
    main()

"""用训练好的权重对单张图片做预测：推理阶段的完整姿势。

训练脚本和推理脚本是两份东西，推理这一侧最容易漏的四件事：
  1. 网络结构必须和训练时**完全一致**，否则 load_state_dict 直接报 key 不匹配；
  2. 必须 model.eval()，否则 Dropout 还在随机失活，同一张图连问两次答案不同；
  3. 必须 torch.no_grad()，推理不需要梯度，省内存也更快；
  4. 预处理必须和训练时**逐字一致**——训练时怎么归一化，这里就怎么归一化。
"""
import os

import torch
import torch.nn as nn

# CIFAR-10 的类别顺序由数据集固定，不能自己重排
CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck']
CN = ['飞机', '汽车', '鸟', '猫', '鹿', '狗', '青蛙', '马', '船', '卡车']

WEIGHT = os.environ.get('CIFAR_WEIGHT', 'model/image_classification.pth')


class ImageClassification(nn.Module):
    """必须和训练脚本里的定义一模一样，连层的属性名都不能改。"""

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, kernel_size=3, stride=1)    # 3×32×32 -> 6×30×30
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)       # -> 6×15×15
        self.conv2 = nn.Conv2d(6, 16, kernel_size=3, stride=1)   # -> 16×13×13
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)       # -> 16×6×6
        self.linear1 = nn.Linear(576, 120)                       # 576 = 16×6×6
        self.linear2 = nn.Linear(120, 84)
        self.out = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv1(x)))
        x = self.pool2(torch.relu(self.conv2(x)))
        x = x.reshape(x.size(0), -1)
        x = torch.relu(self.linear1(x))
        x = torch.relu(self.linear2(x))
        return self.out(x)


def load_model(weight=WEIGHT):
    model = ImageClassification()
    state = torch.load(weight, map_location='cpu')
    model.load_state_dict(state)
    model.eval()                      # 第 2 件事：切到评估模式
    return model


def preprocess(img_hwc):
    """把 (H, W, C)、取值 0~255 的图片变成网络要的 (1, 3, 32, 32)、0.0~1.0。

    这两步必须和训练时的 ToTensor() 完全对齐：
      ToTensor 做的就是「HWC -> CHW」+「除以 255」。
    少做一步归一化，预测结果会错得莫名其妙而且不报错。
    """
    x = torch.as_tensor(img_hwc, dtype=torch.float32)
    if x.max() > 1.5:                 # 还是 0~255 的量纲就归一化
        x = x / 255.0
    x = x.permute(2, 0, 1)            # HWC -> CHW
    return x.unsqueeze(0)             # 补出 batch 维


@torch.no_grad()                      # 第 3 件事：整个函数都不追踪梯度
def predict(model, x, topk=3):
    logits = model(x)                            # 原始分数，不是概率
    probs = torch.softmax(logits, dim=1)[0]      # 推理阶段自己补 softmax
    values, indices = probs.topk(topk)
    return [(CLASSES[i], CN[i], v.item()) for v, i in zip(values, indices)]


def main():
    if not os.path.exists(WEIGHT):
        print('找不到权重文件 %s，先跑训练脚本，或用环境变量 CIFAR_WEIGHT 指定路径' % WEIGHT)
        return

    model = load_model()
    print('参数量', sum(p.numel() for p in model.parameters()))

    # 真实使用时把这里换成 plt.imread('xxx.jpg') 读进来的数组，
    # 尺寸不是 32×32 的要先缩放，否则拉平维度对不上。
    fake = torch.randint(0, 256, (32, 32, 3))
    x = preprocess(fake)
    print('送进网络的形状', tuple(x.shape))

    for name, cn, p in predict(model, x):
        print('  %-12s %-4s %.2f%%' % (name, cn, p * 100))

    # 验证 eval() 确实生效：同一张图连问两次，结果必须完全一致
    a = predict(model, x)
    b = predict(model, x)
    print('两次预测是否完全一致：', a == b)


if __name__ == '__main__':
    main()

"""CNN 骨架模板：改通道数和层数，拉平那一步用代码自动算，别手填。"""
import torch
import torch.nn as nn


class MyCNN(nn.Module):
    def __init__(self, in_channels=3, image_size=32, n_class=10):
        super().__init__()

        # TODO: 按需要增减「卷积 + 激活 + 池化」块；通道数一般逐层翻倍
        self.features = nn.Sequential(
            nn.Conv2d(in_channels, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )

        # 用一次假前向把拉平后的长度问出来，比人肉套公式可靠
        with torch.no_grad():
            dummy = torch.zeros(1, in_channels, image_size, image_size)
            flat_dim = self.features(dummy).flatten(1).shape[1]
        print('拉平后维度 =', flat_dim)

        # TODO: 分类头的宽度按任务调
        self.classifier = nn.Sequential(
            nn.Linear(flat_dim, 128),
            nn.ReLU(),
            nn.Dropout(p=0.3),
            nn.Linear(128, n_class),       # 不加 softmax，交给 CrossEntropyLoss
        )

    def forward(self, x):
        x = self.features(x)
        x = x.flatten(1)                   # 等价 x.reshape(x.size(0), -1)
        return self.classifier(x)


if __name__ == '__main__':
    net = MyCNN()
    print(net(torch.randn(4, 3, 32, 32)).shape)    # torch.Size([4, 10])
    print('参数量', sum(p.numel() for p in net.parameters()))

"""CNN 里五种形状报错的原地复现：每一种都先触发、再打印修法。

全部用 try/except 抓住，跑到底不会中断。读一遍不如自己触发一遍。
"""
import torch
import torch.nn as nn


def case(title):
    print('\n' + '=' * 60)
    print(title)


if __name__ == '__main__':
    torch.manual_seed(0)

    # ------------------------------------------------------------ 一
    case('一、把 (H, W, C) 直接喂进卷积层')
    hwc = torch.randn(32, 32, 3)          # 图像库给出的顺序
    conv = nn.Conv2d(3, 6, kernel_size=3)
    try:
        conv(hwc)
    except RuntimeError as exc:
        print('报错：', str(exc).split('\n')[0])
    fixed = hwc.permute(2, 0, 1).unsqueeze(0)     # HWC -> CHW -> 补 batch 维
    print('修法：permute(2,0,1).unsqueeze(0) ->', tuple(fixed.shape))
    print('输出：', tuple(conv(fixed).shape))

    # ------------------------------------------------------------ 二
    case('二、忘了补 batch 维，单张图片直接送进去')
    chw = torch.randn(3, 32, 32)
    try:
        conv(chw)
    except RuntimeError as exc:
        print('报错：', str(exc).split('\n')[0])
    print('修法：x.unsqueeze(0) ->', tuple(conv(chw.unsqueeze(0)).shape))
    print('说明：网络永远按「一批」处理数据，哪怕这一批只有一张。')

    # ------------------------------------------------------------ 三
    case('三、拉平后的维度写错（最高频的一种）')

    class WrongFlat(nn.Module):
        def __init__(self, flat_dim):
            super().__init__()
            self.conv1 = nn.Conv2d(3, 6, kernel_size=3)
            self.pool1 = nn.MaxPool2d(2, 2)
            self.conv2 = nn.Conv2d(6, 16, kernel_size=3)
            self.pool2 = nn.MaxPool2d(2, 2)
            self.fc = nn.Linear(flat_dim, 120)

        def forward(self, x):
            x = self.pool1(torch.relu(self.conv1(x)))
            x = self.pool2(torch.relu(self.conv2(x)))
            x = x.reshape(x.size(0), -1)
            return self.fc(x)

    x = torch.randn(4, 3, 32, 32)
    try:
        WrongFlat(400)(x)                 # 400 是 5×5 核网络的数，这里该是 576
    except RuntimeError as exc:
        print('报错：', str(exc).split('\n')[0])
    print('修法一：按公式推 16×6×6 = 576')
    print('输出：', tuple(WrongFlat(576)(x).shape))
    print('修法二（更稳）：用一次假前向把这个数问出来，见下面 probe_flat_dim')

    # ------------------------------------------------------------ 四
    case('四、拉平时用了 reshape(-1) 而不是 reshape(batch, -1)')
    feat = torch.randn(4, 16, 6, 6)
    wrong = feat.reshape(-1)              # 把整个 batch 揉成一根长向量
    right = feat.reshape(feat.size(0), -1)
    print('reshape(-1)          ->', tuple(wrong.shape), ' batch 维没了')
    print('reshape(size(0), -1) ->', tuple(right.shape), ' 正确')
    fc = nn.Linear(576, 120)
    try:
        fc(wrong)
    except RuntimeError as exc:
        print('报错：', str(exc).split('\n')[0])
    print('修法：永远保留第 0 维，或者直接用 x.flatten(1)')

    # ------------------------------------------------------------ 五
    case('五、上一层输出通道和下一层输入通道接不上')
    bad = nn.Sequential(
        nn.Conv2d(3, 6, kernel_size=3),
        nn.Conv2d(8, 16, kernel_size=3),   # 上一层出 6 通道，这里却写 8
    )
    try:
        bad(torch.randn(1, 3, 32, 32))
    except RuntimeError as exc:
        print('报错：', str(exc).split('\n')[0])
    print('修法：后一层的 in_channels 必须等于前一层的 out_channels。')

    # ------------------------------------------------------------ 通用工具
    case('通用工具：与其手推，不如让代码自己量一遍')

    def probe_flat_dim(features, in_channels=3, image_size=32):
        """拿一个全 0 的假输入走一遍特征提取部分，量出拉平后的长度。"""
        with torch.no_grad():
            dummy = torch.zeros(1, in_channels, image_size, image_size)
            return features(dummy).flatten(1).shape[1]

    features = nn.Sequential(
        nn.Conv2d(3, 6, kernel_size=3), nn.ReLU(), nn.MaxPool2d(2, 2),
        nn.Conv2d(6, 16, kernel_size=3), nn.ReLU(), nn.MaxPool2d(2, 2),
    )
    print('3×32×32 输入下量出来的拉平维度 =', probe_flat_dim(features))
    print('换成 5×5 的核再量一次：')
    features5 = nn.Sequential(
        nn.Conv2d(3, 6, kernel_size=5), nn.ReLU(), nn.MaxPool2d(2, 2),
        nn.Conv2d(6, 16, kernel_size=5), nn.ReLU(), nn.MaxPool2d(2, 2),
    )
    print('拉平维度 =', probe_flat_dim(features5))
    print('这样改核大小、加层、换输入尺寸都不用重新手推。')

    # ------------------------------------------------------------ 逐层打印
    case('排错习惯：在 forward 里逐层打印形状')

    class Verbose(nn.Module):
        def __init__(self, verbose=True):
            super().__init__()
            self.verbose = verbose
            self.conv1 = nn.Conv2d(3, 6, kernel_size=3)
            self.pool1 = nn.MaxPool2d(2, 2)
            self.conv2 = nn.Conv2d(6, 16, kernel_size=3)
            self.pool2 = nn.MaxPool2d(2, 2)

        def forward(self, x):
            steps = [('输入', x)]
            x = torch.relu(self.conv1(x)); steps.append(('卷积1', x))
            x = self.pool1(x);             steps.append(('池化1', x))
            x = torch.relu(self.conv2(x)); steps.append(('卷积2', x))
            x = self.pool2(x);             steps.append(('池化2', x))
            x = x.flatten(1);              steps.append(('拉平', x))
            if self.verbose:
                for name, t in steps:
                    print('  %-8s %s' % (name, tuple(t.shape)))
            return x

    Verbose()(torch.randn(4, 3, 32, 32))
    print('把这段打印留在开发期，定稿前把 verbose 关掉即可。')