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

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)

二、使用卷积层实现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)

猜你喜欢

转载自blog.csdn.net/qq_37760750/article/details/107769343
今日推荐