---

版本	文档	测试	覆盖率	样式	PyPI	Python	PyTorch	Docker

torchlayers 是一个基于 PyTorch 的库， 提供了 自动推断 torch.nn 层的形状和维度 功能，同时还包含当前 SOTA 架构中使用的额外模块（例如 Efficient-Net）。

上述功能无需用户干预（除了调用一次 torchlayers.build），与 Keras 中的做法类似。

主要功能：

• 对大多数 torch.nn 模块进行 形状推断（包括卷积层、循环层、Transformer、注意力机制和线性层等）
• 维度推断（例如，torchlayers.Conv 可根据输入形状自动选择 torch.nn.Conv1d/2d/3d）
• 自定义模块的形状推断（见示例部分）
• 额外的[类似于 Keras 的]层（如 torchlayers.Reshape 或 torchlayers.StandardNormalNoise）
• 来自 ImageNet 竞赛的 SOTA 层（如 PolyNet、 Squeeze-And-Excitation、 StochasticDepth)
• 实用的默认设置（如 Conv 的 "same" 填充和默认 kernel_size=3，以及 Dropout 率等）
• 零开销且支持 torchscript

请记住，这个库的工作方式几乎与原生 PyTorch 完全一致。这意味着你可以使用 Sequential，利用 torch.nn.Module 定义任意复杂度的网络， 创建具有形状推断功能的新层等等。

请参阅下方内容以更好地理解该库的功能。

示例

如需了解完整功能，请查阅 torchlayers 文档。 以下示例将介绍你需要掌握的所有基本概念。

基础分类器

通过 torchlayers 可以使用 所有 torch.nn 模块，并且 每个带有输入形状的模块 都会被相应地替换为可推断输入的等效版本。

import torchlayers as tl


class Classifier(tl.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = tl.Conv2d(64, kernel_size=6)
        self.conv2 = tl.Conv2d(128, kernel_size=3)
        self.conv3 = tl.Conv2d(256, kernel_size=3, padding=1)
        # 新增层，在下一个示例中会详细介绍
        self.pooling = tl.GlobalMaxPool()
        self.dense = tl.Linear(10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.relu(self.conv3(x))
        return self.dense(self.pooling(x))

# 传递模型及示例输入
clf = tl.build(Classifier(), torch.randn(1, 3, 32, 32))

上述代码中使用了 torchlayers.Linear(out_features=10)。它“等价”于原生 PyTorch 的 torch.nn.Linear(in_features=?, out_features=10)，其中 in_features 将在调用 torchlayers.build 时 根据示例输入自动推断出来。

同样的情况也适用于 torch.nn.Conv2d(in_channels, out_channels, kernel_size, ...)， 可以直接用 tl.Conv2d(out_channels, kernel_size, ...) 替代。

只需记得将示例输入传递给网络即可！

图像与文本分类一体化！

• 我们将使用同一个“模型”来处理这两项任务。首先让我们用 torch.nn 和 torchlayers 来定义它：

import torch
import torchlayers as tl

# torch.nn 和 torchlayers 可以轻松混合使用
model = torch.nn.Sequential(
    tl.Conv(64),  # 只需指定输出通道数
    torch.nn.ReLU(),  # 随意使用 torch.nn
    tl.BatchNorm(),  # 根据输入推断 BatchNormNd
    tl.Conv(128),  # 默认卷积核大小为 3
    tl.ReLU(),
    tl.Conv(256, kernel_size=11),  # 默认采用 "same" 填充
    tl.GlobalMaxPool(),  // 类似 Keras 中的全局池化
    tl.Linear(10),  // 输出 10 个类别的结果
)

print(model)

上述代码会生成如下模型摘要（注意尚未推断出的值用问号表示）：

Sequential(
  (0): Conv(in_channels=?, out_channels=64, kernel_size=3, stride=1, padding=same, dilation=1, groups=1, bias=True, padding_mode=zeros)
  (1): ReLU()
  (2): BatchNorm(num_features=?, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (3): Conv(in_channels=?, out_channels=128, kernel_size=3, stride=1, padding=same, dilation=1, groups=1, bias=True, padding_mode=zeros)
  (4): ReLU()
  (5): Conv(in_channels=?, out_channels=256, kernel_size=11, stride=1, padding=same, dilation=1, groups=1, bias=True, padding_mode=zeros)
  (6): GlobalMaxPool()
  (7): Linear(in_features=?, out_features=10, bias=True)
)

• 现在你可以通过示例输入构建/实例化你的模型（这里以类似 MNIST 的数据为例）：

mnist_model = tl.build(model, torch.randn(1, 3, 28, 28))

• 如果你是在做文本分类，也可以用不同的 input shape 来构建相同的模型（例如，使用 300 维的预训练嵌入进行文本分类）：

```python
# [批次, 嵌入, 时间步], 第一维必须大于1，BatchNorm1d才能正常工作
text_model = tl.build(model, torch.randn(2, 300, 1))

• 最后，在实例化之后，你可以打印这两个模型，如下并排显示以提高可读性（请注意维度的不同，例如 torchlayers.build 后的 Conv2d 和 Conv1d）：

                # 文本分类器                 MNIST分类器

                Sequential(                       Sequential(
                  (0): Conv1d(300, 64)              (0): Conv2d(3, 64)
                  (1): ReLU()                       (1): ReLU()
                  (2): BatchNorm1d(64)              (2): BatchNorm2d(64)
                  (3): Conv1d(64, 128)              (3): Conv2d(64, 128)
                  (4): ReLU()                       (4): ReLU()
                  (5): Conv1d(128, 256)             (5): Conv2d(128, 256)
                  (6): GlobalMaxPool()              (6): GlobalMaxPool()
                  (7): Linear(256, 10)              (7): Linear(256, 10)
                )                                 )

正如你所看到的，这两个模块都被“编译”成了原始的 PyTorch 层。

具有形状推断能力的自定义模块

用户可以定义任意模块，并使用 torchlayers.infer 函数使其具备形状推断能力：

 # 定义了一个带有 in_features 的类
 # 为了与可进行形状推断的版本区分开来，建议使用 _ 前缀和 Impl 后缀
class _MyLinearImpl(torch.nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.randn(out_features, in_features))
        self.bias = torch.nn.Parameter(torch.randn(out_features))

    def forward(self, inputs):
        return torch.nn.functional.linear(inputs, self.weight, self.bias)

MyLinear = tl.infer(_MyLinearImpl)

# 构建并像使用本库中的其他层一样使用
layer =tl.build(MyLinear(out_features=32), torch.randn(1, 64))
layer(torch.randn(1, 64))

默认情况下，在首次 forward 传递过程中，会使用 inputs.shape[1] 作为 in_features 的值。如果你希望使用不同的索引（例如使用 inputs.shape[3] 进行推断），可以使用 MyLayer = tl.infer(_MyLayerImpl, index=3) 作为装饰器。

具有倒残差瓶颈和像素洗牌的自动编码器

如有需要，请查看代码注释和文档。如果你不确定什么是自动编码器，可以参考这篇示例博客文章。

以下是针对类似 ImageNet 图像的卷积去噪自动编码器示例。可以将其视为展示 torchlayers 提供的各种层和构建块功能的一个演示。

# 输入为 3 x 256 x 256，用于 ImageNet 重建
class AutoEncoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = tl.Sequential(
            tl.StandardNormalNoise(),  # 对输入图像添加噪声
            tl.Conv(64, kernel_size=7),
            tl.activations.Swish(),  # 直接访问模块 .activations
            tl.InvertedResidualBottleneck(squeeze_excitation=False),
            tl.AvgPool(),  # 形状为 64 x 128 x 128，默认内核大小为 2
            tl.HardSwish(),  # 可直接通过 tl 访问
            tl.SeparableConv(128),  # 将通道数增加到 128
            tl.InvertedResidualBottleneck(),  # 默认带挤压激励机制
            torch.nn.ReLU(),
            tl.AvgPool(),  # 形状为 128 x 64 x 64，默认内核大小为 2
            tl.DepthwiseConv(256),  # 深度可分离卷积更易于使用
            # 像 PolyNet 一样将输入三次通过相同的权重
            tl.Poly(tl.InvertedResidualBottleneck(), order=3),
            tl.ReLU(),  # 所有 torch.nn 层都可以通过 tl 访问
            tl.MaxPool(),  # 形状为 256 x 32 x 32
            tl.Fire(out_channels=512),  # 形状为 512 x 32 x 32
            tl.SqueezeExcitation(hidden=64),
            tl.InvertedResidualBottleneck(),
            tl.MaxPool(),  # 形状为 512 x 16 x 16
            tl.InvertedResidualBottleneck(squeeze_excitation=False),
            # 随机以 50% 的概率关闭最后两层
            tl.StochasticDepth(
                torch.nn.Sequential(
                    tl.InvertedResidualBottleneck(squeeze_excitation=False),
                    tl.InvertedResidualBottleneck(squeeze_excitation=False),
                ),
                p=0.5,
            ),
            tl.AvgPool(),  # 形状为 512 x 8 x 8
        )

        # 解码器部分更为“标准”
        self.decoder = tl.Sequential(
            tl.Poly(tl.InvertedResidualBottleneck(), order=2),
            # 调用 `build` 后默认具有 ICNR 初始化
            tl.ConvPixelShuffle(out_channels=512, upscale_factor=2),
            # PixelShuffle 后形状为 512 x 16 x 16
            tl.Poly(tl.InvertedResidualBottleneck(), order=3),
            tl.ConvPixelShuffle(out_channels=256, upscale_factor=2),
            # 形状为 256 x 32 x 32
            tl.Poly(tl.InvertedResidualBottleneck(), order=3),
            tl.ConvPixelShuffle(out_channels=128, upscale_factor=2),
            # 形状为 128 x 64 x 64
            tl.Poly(tl.InvertedResidualBottleneck(), order=4),
            tl.ConvPixelShuffle(out_channels=64, upscale_factor=2),
            # 形状为 64 x 128 x 128
            tl.InvertedResidualBottleneck(),
            tl.Conv(256),
            tl.Dropout(),  # 默认为 0.5，且对图像使用 Dropout2d
            tl.Swish(),
            tl.InstanceNorm(),
            tl.ConvPixelShuffle(out_channels=32, upscale_factor=2),
            # 形状为 32 x 256 x 256
            tl.Conv(16),
            tl.Swish(),
            tl.Conv(3),
            # 最终形状为 3 x 256 x 256
        )

    def forward(self, inputs):
        return self.decoder(self.encoder(inputs))

现在可以实例化该模块，并像往常一样使用 torch.nn.MSELoss 来训练它。

autoencoder = tl.build(AutoEncoder(), torch.randn(1, 3, 256, 256))

安装

pip

最新版本：

pip install --user torchlayers

夜间版本：

pip install --user torchlayers-nightly

Docker

在 dockerhub 上提供了 仅 CPU 版本以及各种 支持 GPU 的镜像。

对于 CPU 快速入门，执行以下命令：

docker pull szymonmaszke/torchlayers:18.04

夜间构建版本也可用，只需在标签前加上 nightly_ 即可。如果要使用 GPU 镜像，请确保已安装 nvidia/docker，并正确配置其运行时环境。

贡献

如果您发现任何问题，或者希望添加某些功能（或自行实现），请新建议题或创建拉取请求。

torchlayers 快速上手指南

torchlayers 是一个基于 PyTorch 的库，提供 torch.nn 层的自动形状和维度推断功能，并包含当前 SOTA 架构（如 Efficient-Net）中的额外构建模块。其使用体验类似 Keras，无需手动计算输入维度，只需调用一次 torchlayers.build 即可。

环境准备

在开始之前，请确保您的系统满足以下要求：

• Python: >= 3.7
• PyTorch: >= 1.3.0

安装步骤

您可以使用 pip 直接安装 torchlayers：

pip install torchlayers

如果需要使用国内镜像源加速安装，可以使用以下命令：

pip install torchlayers -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

torchlayers 的核心优势在于自动推断层的大小（如 in_channels 或 in_features）。以下是两种常见的使用模式。

1. 定义自定义模型类

您可以继承 tl.Module 来定义网络，无需指定输入通道数，只需在构建时传入示例输入。

import torch
import torchlayers as tl


class Classifier(tl.Module):
    def __init__(self):
        super().__init__()
        # 无需指定 in_channels，库会自动推断
        self.conv1 = tl.Conv2d(64, kernel_size=6)
        self.conv2 = tl.Conv2d(128, kernel_size=3)
        self.conv3 = tl.Conv2d(256, kernel_size=3, padding=1)
        
        # 全局最大池化
        self.pooling = tl.GlobalMaxPool()
        
        # 无需指定 in_features，库会自动推断
        self.dense = tl.Linear(10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.relu(self.conv3(x))
        return self.dense(self.pooling(x))

# 关键步骤：传入模型实例和任意形状的示例输入进行构建
# 此处示例输入为 batch_size=1, channels=3, height=32, width=32
clf = tl.build(Classifier(), torch.randn(1, 3, 32, 32))

2. 使用 Sequential 快速搭建

torchlayers 可以与原生 torch.nn 模块混合使用。通过 tl.build，同一个模型结构可以根据不同的输入形状编译成不同维度的网络（例如从 1D 文本分类切换到 2D 图像分类）。

import torch
import torchlayers as tl

# 定义模型结构
model = torch.nn.Sequential(
    tl.Conv(64),          # 仅指定输出通道，默认 kernel_size=3, padding="same"
    torch.nn.ReLU(),      # 可以混用原生 PyTorch 层
    tl.BatchNorm(),       # 自动推断为 BatchNorm1d/2d/3d
    tl.Conv(128),         # 默认 kernel_size=3
    tl.ReLU(),
    tl.Conv(256, kernel_size=11), 
    tl.GlobalMaxPool(),   # 类似 Keras 的全局池化
    tl.Linear(10),        # 输出 10 类
)

# 场景 A: 构建用于 MNIST 图像分类的模型 (2D)
mnist_model = tl.build(model, torch.randn(1, 3, 28, 28))

# 场景 B: 构建用于文本分类的模型 (1D)
# 输入形状: [batch, embedding_dim, timesteps]
text_model = tl.build(model, torch.randn(2, 300, 1))

# 打印查看结构，你会发现 Conv 自动变为了 Conv2d 或 Conv1d
print(mnist_model)
print(text_model)

核心提示：

• 所有 torch.nn 模块均可通过 torchlayers 使用。
• 务必在定义好网络结构后，调用 tl.build(model, example_input) 来完成形状推断和初始化。
• 构建后的模型即为标准的 PyTorch 模型，可正常进行训练、保存和导出 TorchScript。