深度可分离卷积的pytorch代码

以下是一个深度可分离卷积的PyTorch代码示例：

import torch.nn as nn

class SeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super(SeparableConv2d, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=in_channels)
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

在这个代码中，我们定义了一个名为`SeparableConv2d`的类，它继承自`nn.Module`。这个类包含两个卷积层，一个是深度卷积（`depthwise`），另一个是逐点卷积（`pointwise`）。深度卷积使用`groups=in_channels`来实现深度可分离卷积，其中每个输入通道都处理为一个单独的输出通道。逐点卷积则将所有深度卷积输出的通道压缩到一个输出通道上。

在`forward`函数中，我们首先对输入进行深度卷积，然后对结果进行逐点卷积，并返回最终结果。

使用这个类非常简单。例如，我们可以在一个`nn.Sequential`模型中添加一个深度可分离卷积层：

model = nn.Sequential(
    nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
    nn.ReLU(inplace=True),
    SeparableConv2d(32, 64, kernel_size=3, stride=1, padding=1),
    nn.ReLU(inplace=True),
    nn.MaxPool2d(kernel_size=2, stride=2),
    SeparableConv2d(64, 128, kernel_size=3, stride=1, padding=1),
    nn.ReLU(inplace=True),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    nn.Linear(128 * 8 * 8, 10)
)