使用转置卷积和卷积分别实现 pytorch 中 nn.PixelShuffle 层的正运算和逆运算。_shuffle层

pytorch 中 nn.PixelShuffle 层主要是将Tensor的通道数降低4倍的同时将其分辨率扩大2倍，但整个过程是不会改变Tensor中的数值的。简单理解就是，nn.PixelShuffle 层的输入和输出是相同的。nn.PixelShuffle 层的具体运算请参见pytorch官方介绍，我主要是使用转置卷积来实现与 nn.PixelShuffle 层相同的功能；并可以使用卷积层实现其相反的功能，也就是它的逆运算。
详细代码如下：
一、使用转置卷积层实现nn.PixelShuffle 层的功能：

# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import torch.nn.functional as F

def pixel_shuffle(input, upsacle_factor):
    '''
        input: (batchSize, c, w, h)
        downscale_factor: k
        (batchSize, c, w, h) -> (batchSize, c/k/k, k*w, k*h)
    '''
    c, w, h = input.shape[1], input.shape[2], input.shape[3]
    kernel = torch.zeros(size = [c, 1, upsacle_factor, upsacle_factor], device=input.device)
    for y in range(upsacle_factor):
        for x in range(upsacle_factor):
            kernel[x + y * upsacle_factor::upsacle_factor * upsacle_factor, 0, y, x] = 1
    # print('kernel:', kernel, kernel.shape)
    return F.conv_transpose2d(input, kernel, stride=upsacle_factor, groups=int(c / upsacle_factor / upsacle_factor))

class Pixel_Shuffle(nn.Module):
    def __init__(self, upsacle_factor):
        super(Pixel_Shuffle, self).__init__()
        self.upsacle_factor = upsacle_factor

    def forward(self, input):
        '''
        input: (batchSize, c, w, h)
        downscale_factor: k
        (batchSize, c, w, h) -> (batchSize, c/k/k, k*w, k*h)
        '''
        return pixel_shuffle(input, self.upsacle_factor)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

二、使用卷积层实现nn.PixelShuffle 层的相反功能：

def pixel_unshuffle(input, downscale_factor):
    '''
    input: batchSize * c * k*w * k*h
    downscale_factor: k
    batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
    '''
    c = input.shape[1]
    kernel = torch.zeros(size = [downscale_factor * downscale_factor * c, 1, downscale_factor, downscale_factor],
                        device = input.device)
    for y in range(downscale_factor):
        for x in range(downscale_factor):
            kernel[x + y * downscale_factor::downscale_factor * downscale_factor, 0, y, x] = 1
    return F.conv2d(input, kernel, stride = downscale_factor, groups = c)

class Pixel_UnShuffle(nn.Module):
    def __init__(self, downscale_factor):
        super(Pixel_UnShuffle, self).__init__()
        self.downscale_factor = downscale_factor

    def forward(self, input):
        '''
        input: batchSize * c * k*w * k*h
        downscale_factor: k
        batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
        '''
        return pixel_unshuffle(input, self.downscale_factor)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26