Hi,
I’m trying to implement SRGAN 31 in pytorch. The paper proposes a pixel shuffle layer in order to upscale images. However, when using such layer, the time per epoch increases drastically, a behavior not observed when running the same code in Tensorflow. It seems that the -pytorch- code is not utilizing the GPU correctly, as the volatile gpu memory oscillates between 10-90% : when removing the layer it stays around 98-99%, as it does with regular gans. Therefore, I’m wondering if the slow behavior is because of my code, or due to pytorch.
Here is my pytorch implementation :
def pixel_shuffle_layer(x, r, n_split):
    def PS(x, r):
        # assumes tf ordering
        bs, a, b, c = x.size()
        x = x.contiguous().view(bs, a, b, r, r)
        x = x.permute(0,1,2,4,3)
        x = torch.chunk(x, a, dim=1)
        x = torch.cat([x_.squeeze() for x_ in x], dim=2)
        x = torch.chunk(x, b, dim=1)
        x = torch.cat([x_.squeeze() for x_ in x], dim=2)
        return x.view(bs, a*r, b*r, 1)

    # put in tf ordering
    x = x.permute(0,2,3,1)
    xc = torch.chunk(x, n_split, dim=3)
    x = torch.cat([PS(x_, r) for x_ in xc], dim=3)
    # put back in th ordering
    x = x.permute(0,3,1,2)
    return x

Here is the tensorflow implementation 18
def pixel_shuffle_layer(x, r, n_split):
    def PS(x, r):
        bs, a, b, c = x.get_shape().as_list()
        x = tf.reshape(x, (bs, a, b, r, r))
        x = tf.transpose(x, (0, 1, 2, 4, 3))
        x = tf.split(x, a, 1)
        x = tf.concat([tf.squeeze(x_) for x_ in x], 2)
        x = tf.split(x, b, 1)
        x = tf.concat([tf.squeeze(x_) for x_ in x], 2)
        return tf.reshape(x, (bs, a*r, b*r, 1))

    xc = tf.split(x, n_split, 3)
    return tf.concat([PS(x_, r) for x_ in xc], 3)

Finally, here is how I’m calling the layer in my forward pass :
def forward(self, input):
    val = self.deconv_one(input)
    val = nn.ReLU()(val)
    shortcut = val 

    for i in range(self.B):
        mid = val
        val = self.blocks[i](val)
        val = val + mid

    val = self.deconv_a(val)
    if self.bn : val = self.bn_a(val)
    val = val + shortcut

    val = self.deconv_b(val)
    bs, C, H, W = val.size()
    val = pixel_shuffle_layer(val, 2, C / 2**2)
    val = nn.ReLU()(val)

    val = self.deconv_c(val)
    bs, C, H, W = val.size()
    val = pixel_shuffle_layer(val, 2, C / 2**2)
    val = nn.ReLU()(val)

    val = self.deconv_d(val)
    if self.last_nl is not None : 
        val = self.last_nl(val)

    return val

Thanks in advance,
Lucas