ROCm上运行网络中的网络（NiN）

7.3. 网络中的网络（NiN） — 动手学深度学习 2.0.0 documentation (d2l.ai)

代码

import torch
from torch import nn
from d2l import torch as d2l


def nin_block(in_channels, out_channels, kernel_size, strides, padding):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size, strides, padding),
        nn.ReLU(),
        nn.Conv2d(out_channels, out_channels, kernel_size=1), nn.ReLU(),
        nn.Conv2d(out_channels, out_channels, kernel_size=1), nn.ReLU())

net = nn.Sequential(
    nin_block(1, 96, kernel_size=11, strides=4, padding=0),
    nn.MaxPool2d(3, stride=2),
    nin_block(96, 256, kernel_size=5, strides=1, padding=2),
    nn.MaxPool2d(3, stride=2),
    nin_block(256, 384, kernel_size=3, strides=1, padding=1),
    nn.MaxPool2d(3, stride=2),
    nn.Dropout(0.5),
    # 标签类别数是10
    nin_block(384, 10, kernel_size=3, strides=1, padding=1),
    nn.AdaptiveAvgPool2d((1, 1)),
    # 将四维的输出转成二维的输出，其形状为(批量大小,10)
    nn.Flatten())

X = torch.rand(size=(1, 1, 224, 224))
for layer in net:
    X = layer(X)
    print(layer.__class__.__name__,'output shape:\t', X.shape)

lr, num_epochs, batch_size = 0.1, 10, 128
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size, resize=224)
d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())

代码解析

这段代码展示了一个网络架构（Neural Network Architecture）的实现，使用的是网络中网络（Network In Network, NIN）模型。首先，代码定义了一个函数 nin_block，这个函数创建了NIN块，这是NIN模型中的基本构件。然后，代码构建了一个完整的NIN模型，使用这些NIN块，并对其进行训练和测试。让我们逐步分析一下：
1. 导入必要的包：

   import torch
   from torch import nn
   from d2l import torch as d2l

   导入了PyTorch的核心库，用于构建网络模型和数据处理，以及`d2l`作为`Dive into Deep Learning`书的辅助工具包，用于加载数据和模型训练等功能。
2. 定义`nin_block`函数：

   def nin_block(in_channels, out_channels, kernel_size, strides, padding):
       return nn.Sequential(
           nn.Conv2d(in_channels, out_channels, kernel_size, strides, padding),
           nn.ReLU(),
           nn.Conv2d(out_channels, out_channels, kernel_size=1), nn.ReLU(),
           nn.Conv2d(out_channels, out_channels, kernel_size=1), nn.ReLU())

   这个函数用于构建NIN块，接受输入通道数（`in_channels`）、输出通道数（`out_channels`）、卷积核大小（`kernel_size`）、步长（`strides`）和填充（`padding`）作为参数。NIN块的特点是在传统卷积层之后再接上两个`1x1`的卷积层，这两个`1x1`卷积层被用来增加网络深度，用于替换传统的全连接层。
3. 构建NIN模型：

   net = nn.Sequential(
       ...
       nn.Flatten())

   使用`nn.Sequential`定义了NIN的网络结构，它由多个`nin_block`和`MaxPool2d`、`Dropout`等层经过堆叠后形成，并且在网络最后使用全局平均池化层`AdaptiveAvgPool2d`代替全连接层，然后将输出扁平化（`Flatten`）以匹配类别标签数量。
4. 打印每层的输出形状：

   X = torch.rand(size=(1, 1, 224, 224))
   for layer in net:
       X = layer(X)
       print(layer.__class__.__name__,'output shape:\t', X.shape)

   这部分代码通过随机初始化输入`X`，然后逐层传递和打印每一层的输出形状，以确保网络结构设计对应的输出形状符合预期。
5. 加载数据并进行训练：

   lr, num_epochs, batch_size = 0.1, 10, 128
   train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size, resize=224)
   d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())

   确定学习率`lr`、训练的轮数`num_epochs`和批量大小`batch_size`，调用`d2l.load_data_fashion_mnist`函数加载Fashion-MNIST数据集，其中，图像输入被调整到了224x224像素。然后，`d2l.train_ch6`函数用来训练模型，它将使用GPU进行训练（如果可用的话）。
综上所述，这段代码是使用NIN构建和训练模型来处理Fashion MNIST数据集的示例。