PyTorch深度可分离卷积：优化模型效率

import torch
from collections import OrderedDict

import torch
import torch.nn as nn

def make_layers(block, no_relu_layers):
    layers = []
    for layer_name, v in block.items():
        if 'pool' in layer_name:
            layer = nn.MaxPool2d(kernel_size=v[0], stride=v[1],
                                    padding=v[2])
            layers.append((layer_name, layer))
        else:
            # 使用深度可分离卷积替换标准卷积
            depthwise_conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[0],
                               kernel_size=v[2], stride=v[3],
                               padding=v[4], groups=v[0])
            pointwise_conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1],
                               kernel_size=1, stride=1,
                               padding=0)
            layers.append((layer_name, depthwise_conv2d))
            layers.append(('pointwise_'+layer_name, pointwise_conv2d))
            if layer_name not in no_relu_layers:
                layers.append(('relu_'+layer_name, nn.ReLU(inplace=True)))

    return nn.Sequential(OrderedDict(layers))