Pytorch实现深度可分类卷积

import torch
import torch.nn as nn
from torchsummary import summary
import time

device = torch.device('cuda')

class Common_Convolution(nn.Module):
    def __init__(self, in_chs, out_chs):
        super(Common_Convolution, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_chs,
            out_channels=out_chs,
            kernel_size=(3, 3),
            stride=(1, 1),
            padding=(1, 1),
            groups=1,
        )
    def forward(self, x):
        out = self.conv(x)
        return out

class Depth_Separable_Convolution(nn.Module):
    def __init__(self, in_chs, out_chs):
        super(Depth_Separable_Convolution, self).__init__()
        self.depthwise_conv = nn.Sequential(
            nn.Conv2d(
                in_channels=in_chs,
                out_channels=in_chs,
                kernel_size=(3, 3),
                stride=(1, 1),
                groups=in_chs,
                padding=(1, 1),
            ),
            nn.BatchNorm2d(num_features=in_chs),
            nn.ReLU(),
        )
        self.pointwise_conv = nn.Sequential(
            nn.Conv2d(
                in_channels=in_chs,
                out_channels=out_chs,
                kernel_size=(1, 1),
                stride=(1, 1),
                groups=1,
                padding=(0, 0),
            ),
            nn.BatchNorm2d(num_features=out_chs),
            nn.ReLU(),
        )
    def forward(self, x):
        x = self.depthwise_conv(x)
        out = self.pointwise_conv(x)
        return out

s1 = time.perf_counter()
common_conv = Common_Convolution(in_chs=3, out_chs=30).to(device)
print(summary(common_conv, input_size=(3, 64, 64)))
e1 = time.perf_counter()

s2 = time.perf_counter()
depthwise_conv = Depth_Separable_Convolution(in_chs=3, out_chs=3).to(device)
print(summary(depthwise_conv, input_size=(3, 64, 64)))
e2 = time.perf_counter()
s3 = time.perf_counter()
pointwise_conv = Depth_Separable_Convolution(in_chs=3, out_chs=30).to(device)
print(summary(pointwise_conv, input_size=(3, 64, 64)))
e3 = time.perf_counter()
print('Common convolution: ', e1 - s1)
print('Depth separable convolution: ', (e2 - s2) + (e3 - s3))