Depthwise-STFT卷积复现

《Depthwise-STFT based separable Convolutional Neural Networks》论文复现。该论文提出了使用 STFT 替换 Depthwise 结构的卷积层，达到提速的效果。这里是未经优化的复现了文中提到的 Depthwise-STFT 卷积操作（卷积操作部分可以优化，预计计算了减少 8 倍），由于未优化，速度并没有得到明显的提升。

但是该文提供的改进思路值得学习。

• 改进思路

• 论文中达到的效果：

Depthwise-STFT Conv

卷积操作

import torch
import numpy as np
from torch import nn
import torch.nn.functional as F
import cv2


class STFTConv(nn.Module):
    def __init__(self, in_c, out_c, kernel_size, stride=1, padding=0):
        super(STFTConv, self).__init__()
        self.in_c = in_c
        self.out_c = out_c
        self.stride = stride
        self.padding = padding
        self.kernel_size = kernel_size

        self.Y = self.define_Y(kernel_size)
        n = kernel_size[0]
        a = 1 / n
        self.V = np.array([[a, 0], [0, a], [a, a], [a, -a]])
        w1_r, w1_i = self.kernel_fn(self.V[0], self.Y)
        w2_r, w2_i = self.kernel_fn(self.V[1], self.Y)
        w3_r, w3_i = self.kernel_fn(self.V[2], self.Y)
        w4_r, w4_i = self.kernel_fn(self.V[3], self.Y)

        w1_r = torch.FloatTensor(w1_r).expand(self.out_c, self.in_c, kernel_size[0], kernel_size[1])
        self.w1_r = nn.Parameter(w1_r, requires_grad=False)
        w1_i = torch.FloatTensor(w1_i).expand(self.out_c, self.in_c, kernel_size[0], kernel_size[1])
        self.w1_i = nn.Parameter(w1_i, requires_grad=False)

        w2_r = torch.FloatTensor(w2_r).expand(self.out_c, self.in_c, kernel_size[0], kernel_size[1])
        self.w2_r = nn.Parameter(w2_r, requires_grad=False)
        w2_i = torch.FloatTensor(w2_i).expand(self.out_c, self.in_c, kernel_size[0], kernel_size[1])
        self.w2_i = nn.Parameter(w2_i, requires_grad=False)

        w3_r = torch.FloatTensor(w3_r).expand(self.out_c, self.in_c, kernel_size[0], kernel_size[1])
        self.w3_r = nn.Parameter(w3_r, requires_grad=False)
        w3_i = torch.FloatTensor(w3_i).expand(self.out_c, self.in_c, kernel_size[0], kernel_size[1])
        self.w3_i = nn.Parameter(w3_i, requires_grad=False)

        w4_r = torch.FloatTensor(w4_r).expand(self.out_c, self.in_c, kernel_size[0], kernel_size[1])
        self.w4_r = nn.Parameter(w4_r, requires_grad=False)
        w4_i = torch.FloatTensor(w4_i).expand(self.out_c, self.in_c, kernel_size[0], kernel_size[1])
        self.w4_i = nn.Parameter(w4_i, requires_grad=False)

    def forward(self, X):
        c1_r = F.conv2d(X, self.w1_r, stride=self.stride, padding=self.padding)
        c1_i = F.conv2d(X, self.w1_i, stride=self.stride, padding=self.padding)

        c2_r = F.conv2d(X, self.w2_r, stride=self.stride, padding=self.padding)
        c2_i = F.conv2d(X, self.w2_i, stride=self.stride, padding=self.padding)

        c3_r = F.conv2d(X, self.w3_r, stride=self.stride, padding=self.padding)
        c3_i = F.conv2d(X, self.w3_i, stride=self.stride, padding=self.padding)

        c4_r = F.conv2d(X, self.w4_r, stride=self.stride, padding=self.padding)
        c4_i = F.conv2d(X, self.w4_i, stride=self.stride, padding=self.padding)
        c = torch.cat((c1_r, c1_i, c2_r, c2_i, c3_r, c3_i, c4_r, c4_i), dim=1)
        return torch.abs(c)

    @staticmethod
    def define_Y(kernel_size):
        assert len(kernel_size) % 2 == 0  # 2D
        w, h = kernel_size
        Y = []
        for i in range(w):
            yi = []
            for j in range(h):
                yi.append([[i], [j]])  # 列向量
            Y.append(yi)
        Y = np.array(Y)  # 从1开始
        return Y

    def kernel_fn(self, v, Y):
        w = v.dot(Y)
        w = np.squeeze(w, axis=2)
        return np.cos(2 * np.pi * w), -np.sin(2 * np.pi * w)


def _get_stft_kernels(size, v):
    assert len(size) % 2 == 0  # 2D
    h, w = size
    Y = []
    for i in range(w):
        yi = []
        for j in range(h):
            yi.append([[i], [j]])  # 列向量
        Y.append(yi)
    Y = np.array(Y) + 1  # 从1开始

    def kernel_fn():
        w = v.dot(Y)
        w = np.squeeze(w, axis=2)
        return np.cos(2 * np.pi * w), -np.sin(2 * np.pi * w)

    return kernel_fn()

Depthwise-STFT

class DepthwiseSTFT(nn.Module):
    def __init__(self, in_c, out_c, kernel_size=(3, 3), stride=1, padding=0):
        super(DepthwiseSTFT, self).__init__()
        self.stft = STFTConv(in_c, in_c, kernel_size, stride=stride, padding=padding)
        self.conv = nn.Conv2d(in_c * 8, out_c, kernel_size=(1, 1), stride=1)

    def forward(self, X):
        t = self.stft(X)
        return self.conv(t)

Block

class Block1(nn.Module):
    def __init__(self, in_c, out_c):
        super(Block1, self).__init__()
        self.stft_3 = DepthwiseSTFT(in_c, in_c, kernel_size=(3, 3), padding=1)
        self.stft_5 = DepthwiseSTFT(in_c, in_c, kernel_size=(5, 5), padding=2)
        self.conv = nn.Conv2d(in_c * 2, out_c, kernel_size=(1, 1))

    def forward(self, X):
        d1 = self.stft_3(X)
        d2 = self.stft_5(X)
        d = torch.cat((d1, d2), dim=1)
        return self.conv(d)


class Block2(nn.Module):
    def __init__(self, in_c, b, out_c):
        super(Block2, self).__init__()
        self.b = b
        self.conv1 = nn.Conv2d(in_c, self.b, kernel_size=(1, 1))
        self.stft_3 = DepthwiseSTFT(self.b, self.b, kernel_size=(3, 3), padding=1)
        self.stft_5 = DepthwiseSTFT(self.b, self.b, kernel_size=(5, 5), padding=2)
        self.conv2 = nn.Conv2d(in_c, out_c, kernel_size=(1, 1))
        self.conv = nn.Conv2d(out_c + self.b * 2, out_c, kernel_size=(1, 1))

    def forward(self, X):
        c1 = self.conv1(X)
        c2 = self.conv2(X)
        d1 = self.stft_3(c1)
        d2 = self.stft_5(c1)
        d = torch.cat((d1, d2, c2), dim=1)

        return self.conv(d)

网络整体架构