Mạng Nơ-ron Tích Chập

Mạng nơ-ron tích chập học các hệ phân cấp đặc trưng không gian trực tiếp từ dữ liệu pixel, thay thế các bộ lọc thiết kế thủ công bằng các bộ lọc tối ưu qua gradient. File này đề cập đến cơ chế tích chập, gộp (pooling), bước trượt (stride), giãn nở (dilation), trường tiếp nhận (receptive field), và các kiến trúc mang tính bước ngoặt (LeNet, AlexNet, VGG, ResNet, Inception, EfficientNet) đã định nghĩa nên phân loại ảnh.

• Trong file 01, chúng ta đã thiết kế thủ công các bộ lọc cho phát hiện biên, làm mờ và phát hiện góc. Câu hỏi tự nhiên là: liệu chúng ta có thể học các bộ lọc tối ưu từ dữ liệu không? Đó chính xác là những gì mạng nơ-ron tích chập (CNN) làm.

• Thay vì chọn trọng số bộ lọc bằng tay, CNN học chúng qua hạ gradient (gradient descent, chương 06), khám phá các đặc trưng hữu ích trực tiếp cho tác vụ cần giải quyết.

• Trong chương 06, chúng ta đã giới thiệu phép toán tích chập, kiến thức cơ bản về CNN, và ý tưởng học bộ lọc. Ở đây chúng ta đi sâu hơn vào các đổi mới kiến trúc đã biến CNN trở thành mô hình thống trị trong thị giác máy tính suốt hơn một thập kỷ.

• Nhắc lại phép tích chập cốt lõi: một bộ lọc \(K\) kích thước \(k \times k\) trượt qua feature map đầu vào, tính tích vô hướng tại mỗi vị trí (chương 06). Kích thước đầu ra được kiểm soát bởi ba siêu tham số:

  • Bước trượt (stride): số pixel bộ lọc di chuyển giữa các vị trí. Stride 1 nghĩa là bộ lọc dịch một pixel mỗi lần. Stride 2 nghĩa là dịch hai pixel, giảm một nửa kích thước không gian. Tích chập có stride là một giải pháp thay thế cho pooling để giảm mẫu (downsampling).
  • Đệm (padding): thêm số 0 xung quanh biên đầu vào. Đệm "Same" (\(p = \lfloor k/2 \rfloor\)) giữ nguyên kích thước không gian. Đệm "Valid" (\(p = 0\)) làm giảm chúng.
  • Giãn nở (dilation): chèn khoảng trống giữa các phần tử bộ lọc. Bộ lọc 3x3 với dilation 2 phủ một trường tiếp nhận 5x5 chỉ dùng 9 tham số. Tích chập giãn nở mở rộng trường tiếp nhận mà không tăng khối lượng tính toán.

• Kích thước không gian đầu ra sau tích chập:

\[ \text{out} = \left\lfloor \frac{\text{in} - k + 2p}{s} \right\rfloor + 1 \]

• trong đó \(\text{in}\) là kích thước đầu vào, \(k\) là kích thước kernel, \(p\) là padding, và \(s\) là stride. Công thức này áp dụng độc lập cho chiều cao và chiều rộng.

• Trường tiếp nhận (receptive field) của một nơ-ron là vùng của đầu vào gốc có thể ảnh hưởng đến giá trị của nó.

  • Các tầng đầu có trường tiếp nhận nhỏ (chúng thấy các mẫu cục bộ như đường biên).
  • Các tầng sâu hơn có trường tiếp nhận lớn hơn (chúng thấy các cấu trúc lớn hơn như các bộ phận vật thể).

• Trường tiếp nhận tăng dần theo mỗi tầng: khoảng \(k - 1\) pixel mỗi tầng tích chập (nhiều hơn với stride hoặc dilation).

• Tầng gộp (pooling) giảm kích thước không gian trong khi giữ lại thông tin quan trọng nhất.

  • Max pooling lấy giá trị lớn nhất trong mỗi cửa sổ, giữ lại kích hoạt mạnh nhất (đặc trưng nổi bật nhất).
  • Average pooling lấy giá trị trung bình, làm mượt feature map. Gộp 2x2 với stride 2 giảm một nửa cả hai chiều không gian.

• Gộp trung bình toàn cục (Global Average Pooling - GAP) lấy trung bình toàn bộ phạm vi không gian của mỗi kênh thành một con số duy nhất, tạo ra một vector có độ dài bằng số kênh. GAP thay thế các tầng fully connected ở cuối nhiều kiến trúc hiện đại, giảm mạnh số tham số và hoạt động như một bộ chính quy hoá cấu trúc.

• Chuẩn hoá batch (Batch Normalisation - BatchNorm) chuẩn hoá các kích hoạt trong mỗi mini-batch để có trung bình bằng 0 và phương sai đơn vị, sau đó áp dụng một tỷ lệ và độ dịch có thể học (chương 06). Trong CNN, BatchNorm được áp dụng theo từng kênh: thống kê được tính trên batch và các chiều không gian cho mỗi kênh độc lập. Nó ổn định quá trình huấn luyện, cho phép tốc độ học cao hơn, và hoạt động như một bộ chính quy hoá nhẹ.

• Dropout (chương 06) ngẫu nhiên đặt về 0 các nơ-ron trong quá trình huấn luyện.

• Trong CNN, spatial dropout (Dropout2D) loại bỏ toàn bộ các kênh feature map thay vì từng pixel riêng lẻ, hiệu quả hơn vì các pixel lân cận trong một feature map có tương quan cao.

• Tăng cường dữ liệu (data augmentation) mở rộng tập huấn luyện một cách nhân tạo bằng cách áp dụng các biến đổi ngẫu nhiên lên mỗi ảnh trong quá trình huấn luyện: lật ngang, cắt ngẫu nhiên, xoay, nhiễu màu (điều chỉnh độ sáng, độ tương phản, độ bão hoà, sắc thái), và cutout (che các patch chữ nhật ngẫu nhiên). Mạng nhìn thấy mỗi ảnh dưới nhiều dạng khác nhau, buộc nó học các đặc trưng bất biến với phép biến đổi thay vì ghi nhớ các mẫu pixel cụ thể.

• Các chiến lược tăng cường nâng cao bao gồm Mixup (pha trộn hai ảnh và nhãn của chúng: \(\tilde{x} = \lambda x_i + (1-\lambda) x_j\), \(\tilde{y} = \lambda y_i + (1-\lambda) y_j\)), CutMix (dán một patch chữ nhật từ ảnh này lên ảnh khác và trộn nhãn theo tỷ lệ diện tích), và RandAugment (lấy mẫu ngẫu nhiên một chuỗi các phép tăng cường từ một tập cố định với một tham số cường độ duy nhất).

• Lịch sử các kiến trúc CNN là một câu chuyện về các thiết kế ngày càng sâu hơn, hiệu quả hơn, mỗi thiết kế giải quyết một vấn đề đã hạn chế người tiền nhiệm.

• LeNet-5 (LeCun và cộng sự, 1998) là CNN nguyên bản, được thiết kế cho nhận dạng chữ số viết tay. Hai tầng tích chập tiếp theo là ba tầng fully connected, với average pooling và kích hoạt tanh. Nó chứng minh rằng các bộ lọc học được vượt trội so với đặc trưng thiết kế thủ công, nhưng rất nhỏ so với tiêu chuẩn hiện đại (60K tham số).

• AlexNet (Krizhevsky và cộng sự, 2012) thắng cuộc thi ImageNet với cách biệt lớn, khởi động cuộc cách mạng học sâu. Các đổi mới chính: kích hoạt ReLU (thay vì tanh vốn bị vanishing gradient), dropout để chính quy hoá, tăng cường dữ liệu, và huấn luyện trên GPU. Năm tầng tích chập, ba tầng fully connected, 60 triệu tham số.

• VGG (Simonyan và Zisserman, 2014) cho thấy rằng chỉ dùng các bộ lọc 3x3 xếp chồng sâu hoạt động tốt hơn các bộ lọc lớn hơn. Hai bộ lọc 3x3 xếp chồng có cùng trường tiếp nhận như một bộ lọc 5x5 nhưng ít tham số hơn (\(2 \times 3^2 = 18\) so với \(5^2 = 25\)) và thêm một phi tuyến. VGG-16 (16 tầng) và VGG-19 (19 tầng) vẫn được dùng rộng rãi làm bộ trích xuất đặc trưng. Kiến trúc cực kỳ đơn giản: các khối tích chập với số kênh tăng dần (64, 128, 256, 512), mỗi khối tiếp theo là max pooling.

• GoogLeNet/Inception (Szegedy và cộng sự, 2014) giới thiệu khối Inception: thay vì chọn một kích thước bộ lọc duy nhất, sử dụng song song các tích chập 1x1, 3x3, và 5x5, nối các đầu ra của chúng, và để mạng quyết định tỷ lệ nào hữu ích nhất. Tích chập 1x1 được dùng làm nút cổ chai (bottleneck) trước các bộ lọc lớn hơn để giảm khối lượng tính toán. GoogLeNet đạt độ chính xác cao hơn VGG với ít tham số hơn 12 lần (6.8M so với 138M).

• Khối Inception nắm bắt các đặc trưng ở nhiều tỷ lệ đồng thời. Bộ lọc 1x1 nắm bắt các mẫu điểm, 3x3 nắm bắt kết cấu cục bộ, và 5x5 nắm bắt các cấu trúc lớn hơn. Việc nối kết hợp tất cả các góc nhìn thành một biểu diễn phong phú.

• ResNet (He và cộng sự, 2016) giải quyết vấn đề suy thoái (degradation problem): các mạng sâu hơn hoạt động kém hơn mạng nông hơn, không phải vì overfitting, mà vì chúng khó tối ưu hơn. Giải pháp là kết nối tắt (skip connection / residual connection):

\[ \text{output} = F(x) + x \]

• Tầng học phần dư \(F(x) = \text{output} - x\). Nếu phép biến đổi tối ưu gần với đồng nhất (điều phổ biến trong mạng sâu), việc học một phần dư gần bằng 0 dễ hơn nhiều so với học toàn bộ ánh xạ. Skip connections cũng cung cấp một đường cao tốc gradient trực tiếp, giảm vanishing gradient. ResNet huấn luyện được mạng với 152 tầng, sâu hơn nhiều so với bất cứ thứ gì trước đó.

• Khi kích thước đầu vào và đầu ra khác nhau (do stride hoặc thay đổi kênh), một projection shortcut áp dụng tích chập 1x1 lên \(x\) để khớp kích thước: \(\text{output} = F(x) + W_s x\).

• Khối bottleneck (dùng trong ResNet-50 và sâu hơn) sử dụng ba tích chập: 1x1 để giảm kênh, 3x3 để xử lý không gian, và 1x1 để mở rộng kênh trở lại. Cách này rẻ hơn hai tích chập 3x3 và cho phép các mạng sâu hơn nhiều.

• DenseNet (Huang và cộng sự, 2017) đưa ý tưởng kết nối tắt đi xa hơn: mọi tầng được kết nối với mọi tầng tiếp theo trong một khối dày đặc. Tầng \(l\) nhận các feature map từ tất cả các tầng trước đó làm đầu vào: \(x_l = H_l([x_0, x_1, \ldots, x_{l-1}])\), trong đó \([\cdot]\) biểu thị sự nối dọc theo chiều kênh. Điều này khuyến khích tái sử dụng đặc trưng, tăng cường luồng gradient, và giảm tổng số tham số.

• Các kiến trúc hiệu quả nhắm đến triển khai trên thiết bị di động và phần cứng biên, nơi tài nguyên tính toán, bộ nhớ và năng lượng bị giới hạn.

• MobileNet (Howard và cộng sự, 2017) thay thế tích chập tiêu chuẩn bằng tích chập tách biệt theo chiều sâu (depthwise separable convolution), phân tách phép toán thành hai bước:

  1. Tích chập theo chiều sâu (depthwise convolution): áp dụng một bộ lọc \(k \times k\) duy nhất cho mỗi kênh đầu vào (không có tương tác chéo kênh)
  2. Tích chập điểm (pointwise convolution): áp dụng tích chập 1x1 để kết hợp thông tin qua các kênh

• Một tích chập \(k \times k\) tiêu chuẩn với \(C_{\text{in}}\) kênh đầu vào và \(C_{\text{out}}\) kênh đầu ra tốn \(k^2 \cdot C_{\text{in}} \cdot C_{\text{out}}\) phép nhân mỗi vị trí không gian. Tích chập tách biệt theo chiều sâu tốn \(k^2 \cdot C_{\text{in}} + C_{\text{in}} \cdot C_{\text{out}}\), giảm khoảng \(k^2\) lần. Với bộ lọc 3x3, điều này rẻ hơn khoảng 9 lần.

• MobileNet-V2 giới thiệu khối dư đảo ngược (inverted residual block): mở rộng kênh với tích chập 1x1, áp dụng tích chập theo chiều sâu trong không gian mở rộng, sau đó chiếu ngược lại với tích chập 1x1. Kết nối tắt được đặt trên các tầng hẹp (bottleneck), đảo ngược mẫu ResNet. Tỷ lệ mở rộng thường là 6.

• EfficientNet (Tan và Le, 2019) giới thiệu tỷ lệ hợp chất (compound scaling): thay vì chỉ thay đổi độ sâu, chỉ chiều rộng, hoặc chỉ độ phân giải một cách độc lập, hãy thay đổi cả ba chiều cùng nhau sử dụng một tỷ lệ cố định. Với hệ số tỷ lệ \(\phi\):

\[ \text{depth}: d = \alpha^\phi, \quad \text{width}: w = \beta^\phi, \quad \text{resolution}: r = \gamma^\phi \]

• với điều kiện \(\alpha \cdot \beta^2 \cdot \gamma^2 \approx 2\) (sao cho tổng tính toán tăng gấp đôi mỗi lần \(\phi\) tăng một đơn vị). Một tìm kiếm lưới tìm ra \(\alpha = 1.2\), \(\beta = 1.1\), \(\gamma = 1.15\) làm các tỷ lệ cơ sở. EfficientNet-B0 đến B7 mở rộng dần dần, đạt độ chính xác hiện đại với ít tham số và FLOP hơn nhiều so với các mô hình trước đó.

• ShuffleNet giảm chi phí của tích chập 1x1 (vốn chiếm ưu thế trong các kiến trúc kiểu MobileNet) bằng cách sử dụng tích chập nhóm (group convolution) tiếp theo là xáo trộn kênh (channel shuffle). Tích chập nhóm chia các kênh thành các nhóm và tích chập trong mỗi nhóm độc lập, nhưng điều này ngăn luồng thông tin giữa các nhóm. Phép xáo trộn sắp xếp lại các kênh giữa các nhóm, khôi phục sự pha trộn thông tin với chi phí không đáng kể.

• Học chuyển giao (transfer learning) là thực hành lấy một mô hình đã huấn luyện trên một tác vụ và thích ứng nó với một tác vụ khác. Trong thị giác máy tính, điều này hầu như luôn có nghĩa là bắt đầu từ một mô hình đã tiền huấn luyện trên ImageNet (1.4 triệu ảnh, 1,000 lớp) và thích ứng với một tập dữ liệu chuyên ngành (ảnh y tế, ảnh vệ tinh, khuyết tật sản xuất).

• Trích xuất đặc trưng (feature extraction): đóng băng tất cả các tầng tích chập, loại bỏ phần đầu phân loại cuối cùng, và chỉ huấn luyện một đầu mới ở trên. Các tầng bị đóng băng hoạt động như một bộ trích xuất đặc trưng đa năng. Cách này hiệu quả khi miền đích tương tự ImageNet và tập dữ liệu đích nhỏ.

• Tinh chỉnh (fine-tuning): mở khoá một số hoặc tất cả các tầng tích chập và huấn luyện với tốc độ học nhỏ. Các trọng số tiền huấn luyện đóng vai trò là điểm khởi đầu thay vì các đặc trưng cố định. Tinh chỉnh thường bắt đầu bằng cách mở khoá chỉ các tầng sau (nắm bắt các đặc trưng cấp cao, đặc thù cho tác vụ) và tuỳ chọn mở khoá thêm các tầng sớm hơn.

• Học chuyển giao hiệu quả vì các tầng đầu của CNN học các đặc trưng phổ quát (biên, kết cấu, màu sắc) hữu ích cho mọi tác vụ, trong khi các tầng sau học các đặc trưng cụ thể cho tác vụ. Một mạng được huấn luyện để phân loại động vật vẫn có các bộ phát hiện biên hữu ích cho việc phân loại toà nhà.

• Trực quan hoá CNN tiết lộ mạng đã học được gì và giúp gỡ lỗi hành vi bất ngờ.

• Bản đồ kích hoạt (activation maps / feature maps) hiển thị đầu ra của mỗi bộ lọc cho một ảnh đầu vào nhất định. Kích hoạt tầng đầu trông giống như bản đồ biên; các tầng sâu hơn tạo ra các kích hoạt trừu tượng hơn, thô hơn về mặt không gian.

• Grad-CAM (Gradient-weighted Class Activation Mapping, Selvaraju và cộng sự, 2017) làm nổi bật các vùng của ảnh đầu vào quan trọng nhất cho dự đoán của mô hình. Nó hoạt động bằng cách:

  1. Tính gradient của điểm số lớp mục tiêu đối với feature map của tầng tích chập cuối cùng (sử dụng quy tắc dẫn hàm hàm hợp từ chương 03)
  2. Gộp trung bình toàn cục các gradient này để lấy trọng số quan trọng cho mỗi kênh
  3. Tính tổ hợp có trọng số của các feature map và áp dụng ReLU

\[ L_{\text{Grad-CAM}} = \text{ReLU}\!\left(\sum_k \alpha_k A^k\right), \quad \alpha_k = \frac{1}{Z} \sum_i \sum_j \frac{\partial y^c}{\partial A^k_{ij}} \]

• trong đó \(A^k\) là feature map thứ \(k\), \(\alpha_k\) là trọng số quan trọng cho kênh \(k\), và \(y^c\) là điểm số cho lớp \(c\). Kết quả là một bản đồ nhiệt thô cho thấy vùng nào đã dẫn dắt việc phân loại. ReLU được áp dụng vì chúng ta quan tâm đến các đặc trưng có ảnh hưởng tích cực đến lớp.

• Feature inversion tái tạo một ảnh đầu vào từ biểu diễn đặc trưng của nó bằng cách tối ưu một ảnh ngẫu nhiên để khớp với các đặc trưng mục tiêu (sử dụng hạ gradient trên các giá trị pixel). Điều này tiết lộ thông tin nào mạng giữ lại ở mỗi tầng. Các tầng đầu tái tạo ảnh gần như hoàn hảo; các tầng sâu hơn tạo ra ảnh có thể nhận dạng nhưng bị biến dạng, cho thấy chi tiết không gian tinh tế bị mất trong khi nội dung ngữ nghĩa được bảo toàn.

• Deep Dream và neural style transfer là các ứng dụng sáng tạo của trực quan hoá đặc trưng. Deep Dream tối đa hoá kích hoạt của các nơ-ron tại một tầng được chọn để tạo ra các ảnh siêu thực với mẫu được khuếch đại. Neural style transfer tối ưu một ảnh mục tiêu để khớp với các đặc trưng nội dung (từ tầng sâu) của một ảnh và các đặc trưng phong cách (ma trận Gram của kích hoạt bộ lọc, nắm bắt thống kê kết cấu) của một ảnh khác.

Bài tập lập trình (dùng CoLab hoặc notebook)

1. Cài đặt một CNN đơn giản từ đầu trong JAX với hai tầng tích chập, max pooling và một đầu phân loại. Huấn luyện nó trên một tác vụ phân loại mẫu 2D tổng hợp.

   import jax
   import jax.numpy as jnp
   import jax.lax as lax
   import matplotlib.pyplot as plt
   
   def conv2d(x, kernel, stride=1):
       """Simple 2D convolution for single input, single filter."""
       return lax.conv(x[None, None], kernel[None, None], (stride, stride), 'SAME')[0, 0]
   
   def max_pool(x, size=2):
       """2x2 max pooling."""
       H, W = x.shape
       x = x[:H//size*size, :W//size*size]
       return x.reshape(H//size, size, W//size, size).max(axis=(1, 3))
   
   def init_cnn(key):
       k1, k2, k3 = jax.random.split(key, 3)
       return {
           'conv1': jax.random.normal(k1, (5, 5)) * 0.3,
           'conv2': jax.random.normal(k2, (3, 3)) * 0.3,
           'fc_w': jax.random.normal(k3, (64, 1)) * 0.1,
           'fc_b': jnp.zeros(1),
       }
   
   def forward_cnn(params, img):
       # Conv1 -> ReLU -> Pool
       h = jnp.maximum(0, conv2d(img, params['conv1']))
       h = max_pool(h)
       # Conv2 -> ReLU -> Pool
       h = jnp.maximum(0, conv2d(h, params['conv2']))
       h = max_pool(h)
       # Flatten and classify
       flat = h.ravel()
       # Pad or truncate to fixed size
       flat = jnp.pad(flat, (0, max(0, 64 - len(flat))))[:64]
       logit = (flat @ params['fc_w'] + params['fc_b']).squeeze()
       return jax.nn.sigmoid(logit)
   
   # Generate synthetic data: class 0 = low-freq pattern, class 1 = high-freq
   def make_data(key, n=200):
       images, labels = [], []
       for i in range(n):
           k1, key = jax.random.split(key)
           x, y = jnp.meshgrid(jnp.linspace(0, 4*jnp.pi, 32), jnp.linspace(0, 4*jnp.pi, 32))
           if i < n // 2:
               img = jnp.sin(x) + jax.random.normal(k1, (32, 32)) * 0.1
               labels.append(0)
           else:
               img = jnp.sin(4 * x) * jnp.sin(4 * y) + jax.random.normal(k1, (32, 32)) * 0.1
               labels.append(1)
           images.append(img)
       return images, jnp.array(labels, dtype=jnp.float32)
   
   key = jax.random.PRNGKey(42)
   images, labels = make_data(key)
   params = init_cnn(jax.random.PRNGKey(0))
   
   def loss_fn(params, img, label):
       pred = forward_cnn(params, img)
       return -(label * jnp.log(pred + 1e-7) + (1 - label) * jnp.log(1 - pred + 1e-7))
   
   grad_fn = jax.grad(loss_fn)
   lr = 0.01
   
   for epoch in range(5):
       total_loss = 0.0
       for img, label in zip(images, labels):
           grads = grad_fn(params, img, label)
           params = {k: params[k] - lr * grads[k] for k in params}
           total_loss += loss_fn(params, img, label)
       print(f"Epoch {epoch}: loss = {total_loss / len(images):.4f}")
   
   # Test accuracy
   preds = jnp.array([forward_cnn(params, img) > 0.5 for img in images])
   acc = jnp.mean(preds == labels)
   print(f"Accuracy: {acc:.2%}")
   

2. Trực quan hoá cách các kích thước bộ lọc khác nhau ảnh hưởng đến trường tiếp nhận. Cho thấy hai bộ lọc 3x3 xếp chồng phủ cùng trường tiếp nhận như một bộ lọc 5x5 nhưng với ít tham số hơn.

   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   def compute_receptive_field(layers):
       """Compute receptive field size from a list of (kernel_size, stride) tuples."""
       rf = 1  # start with 1 pixel
       stride_product = 1
       for k, s in layers:
           rf += (k - 1) * stride_product
           stride_product *= s
       return rf
   
   # Compare architectures
   configs = {
       'Single 5x5': [(5, 1)],
       'Two 3x3':    [(3, 1), (3, 1)],
       'Three 3x3':  [(3, 1), (3, 1), (3, 1)],
       'Single 7x7': [(7, 1)],
       '3x3 stride 2 + 3x3': [(3, 2), (3, 1)],
   }
   
   print(f"{'Config':<25} {'RF':>4} {'Params (per channel)':>20}")
   print('-' * 55)
   for name, layers in configs.items():
       rf = compute_receptive_field(layers)
       # Parameters: sum of k^2 for each layer (per input-output channel pair)
       params = sum(k * k for k, s in layers)
       print(f"{name:<25} {rf:>4} {params:>20}")
   
   # Visualise receptive fields
   fig, axes = plt.subplots(1, 3, figsize=(14, 4))
   for ax, (name, rf_size) in zip(axes, [('5x5 filter', 5), ('Two 3x3 filters', 5), ('Three 3x3 filters', 7)]):
       grid = jnp.zeros((9, 9))
       c = 4  # centre
       half = rf_size // 2
       grid = grid.at[c-half:c+half+1, c-half:c+half+1].set(1.0)
       ax.imshow(grid, cmap='Blues', vmin=0, vmax=1)
       ax.set_title(f'{name}\nRF = {rf_size}x{rf_size}')
       ax.set_xticks(range(9)); ax.set_yticks(range(9))
       ax.grid(True, alpha=0.3)
   plt.suptitle('Receptive Field Comparison')
   plt.tight_layout(); plt.show()
   

3. Cài đặt Grad-CAM từ đầu. Cho một CNN đơn giản được xây dựng sẵn, tính bản đồ kích hoạt có trọng số gradient cho một lớp cụ thể và trực quan hoá nó dưới dạng bản đồ nhiệt.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   def simple_cnn(params, img):
       """Simple CNN that returns both the prediction and last conv activations."""
       # Conv layer (our "last conv layer" for Grad-CAM)
       H, W = img.shape
       k = params['conv'].shape[0]
       pad = k // 2
       img_pad = jnp.pad(img, pad, mode='edge')
       activation_map = jnp.zeros((H, W))
       for i in range(H):
           for j in range(W):
               activation_map = activation_map.at[i, j].set(
                   jnp.sum(img_pad[i:i+k, j:j+k] * params['conv'])
               )
       activation_map = jnp.maximum(0, activation_map)  # ReLU
   
       # Global average pool -> dense -> output
       pooled = activation_map.mean()
       logit = pooled * params['w'] + params['b']
       return jax.nn.sigmoid(logit), activation_map
   
   # Create test image: bright region on the left (class indicator)
   img = jnp.zeros((32, 32))
   img = img.at[8:24, 4:16].set(1.0)
   img = img.at[5:10, 20:28].set(0.3)
   
   key = jax.random.PRNGKey(42)
   params = {
       'conv': jax.random.normal(key, (5, 5)) * 0.3,
       'w': jnp.array(2.0),
       'b': jnp.array(-0.5),
   }
   
   # Compute Grad-CAM
   def class_score(params, img):
       pred, _ = simple_cnn(params, img)
       return pred
   
   # Get activation map and gradients
   pred, act_map = simple_cnn(params, img)
   grad_fn = jax.grad(lambda img: simple_cnn(params, img)[0])
   img_grad = grad_fn(img)
   
   # Weight = global average of gradients (simplified 1-channel Grad-CAM)
   alpha = img_grad.mean()
   grad_cam = jnp.maximum(0, alpha * act_map)  # ReLU
   grad_cam = (grad_cam - grad_cam.min()) / (grad_cam.max() - grad_cam.min() + 1e-8)
   
   fig, axes = plt.subplots(1, 3, figsize=(14, 4))
   axes[0].imshow(img, cmap='gray'); axes[0].set_title('Input Image'); axes[0].axis('off')
   axes[1].imshow(act_map, cmap='viridis'); axes[1].set_title('Activation Map'); axes[1].axis('off')
   axes[2].imshow(img, cmap='gray', alpha=0.6)
   axes[2].imshow(grad_cam, cmap='jet', alpha=0.4)
   axes[2].set_title(f'Grad-CAM (pred={pred:.2f})'); axes[2].axis('off')
   plt.tight_layout(); plt.show()
   

4. So sánh tích chập tách biệt theo chiều sâu với tích chập tiêu chuẩn. Đếm số tham số và FLOPs cho cả hai và cho thấy chúng tạo ra đầu ra tương tự với ít tính toán hơn nhiều.

   import jax
   import jax.numpy as jnp
   
   def standard_conv(x, kernel):
       """Standard convolution: (H, W, C_in) * (k, k, C_in, C_out) -> (H, W, C_out)."""
       H, W, C_in = x.shape
       k, _, _, C_out = kernel.shape
       pad = k // 2
       x_pad = jnp.pad(x, ((pad, pad), (pad, pad), (0, 0)), mode='constant')
       out = jnp.zeros((H, W, C_out))
       for i in range(H):
           for j in range(W):
               patch = x_pad[i:i+k, j:j+k, :]  # (k, k, C_in)
               for c in range(C_out):
                   out = out.at[i, j, c].set(jnp.sum(patch * kernel[:, :, :, c]))
       return out
   
   def depthwise_separable_conv(x, dw_kernel, pw_kernel):
       """Depthwise separable: depthwise (k,k,C_in) then pointwise (C_in, C_out)."""
       H, W, C_in = x.shape
       k = dw_kernel.shape[0]
       pad = k // 2
       x_pad = jnp.pad(x, ((pad, pad), (pad, pad), (0, 0)), mode='constant')
   
       # Depthwise: one filter per channel
       dw_out = jnp.zeros((H, W, C_in))
       for i in range(H):
           for j in range(W):
               for c in range(C_in):
                   patch = x_pad[i:i+k, j:j+k, c]
                   dw_out = dw_out.at[i, j, c].set(jnp.sum(patch * dw_kernel[:, :, c]))
   
       # Pointwise: 1x1 conv across channels
       out = dw_out @ pw_kernel
       return out
   
   # Setup
   H, W, C_in, C_out, k = 8, 8, 16, 32, 3
   key = jax.random.PRNGKey(42)
   k1, k2, k3, k4 = jax.random.split(key, 4)
   
   x = jax.random.normal(k1, (H, W, C_in))
   std_kernel = jax.random.normal(k2, (k, k, C_in, C_out)) * 0.1
   dw_kernel = jax.random.normal(k3, (k, k, C_in)) * 0.1
   pw_kernel = jax.random.normal(k4, (C_in, C_out)) * 0.1
   
   # Compare
   std_params = k * k * C_in * C_out
   dw_params = k * k * C_in + C_in * C_out
   
   std_flops = H * W * k * k * C_in * C_out
   dw_flops = H * W * (k * k * C_in + C_in * C_out)
   
   print(f"Standard conv:            {std_params:>8,} params,  {std_flops:>10,} FLOPs")
   print(f"Depthwise separable conv: {dw_params:>8,} params,  {dw_flops:>10,} FLOPs")
   print(f"Parameter reduction:      {std_params / dw_params:.1f}x")
   print(f"FLOP reduction:           {std_flops / dw_flops:.1f}x")
   
   std_out = standard_conv(x, std_kernel)
   ds_out = depthwise_separable_conv(x, dw_kernel, pw_kernel)
   print(f"\nStandard output shape:    {std_out.shape}")
   print(f"Depthwise sep output shape: {ds_out.shape}")