我在尝试将一个torch训练的模型移植到tensorflow（TF）时注意到了一些有趣的事情。当PixelShuffle操作的输出通道大于1时，TF中对应的depth_to_space函数是不等价的（注意：我将TF的输入转换为NHWC，输出转换回NCHW）。我想知道这是预期的行为还是存在误解？
具体地说，

# Torch
torch_out = nn.PixelShuffle(2)(input)

字符串
和

# TF/Keras
input = np.transpose(input, (0, 2, 3, 1)) // Convert to NHWC
keras_input = keras.layers.Input(shape=input.shape[1:])
keras_d2s = keras.layers.Lambda(lambda x: tf.nn.depth_to_space(x, 2))(input)
...
keras_out = np.transpose(keras_d2s, (0, 3, 1, 2)) // Convert back to NCHW

型
和

keras_out != torch_out

型
下面是一个测试平台：

import numpy as np

import torch
import tensorflow as tf

from torch import nn
from tensorflow import keras

class Shuffle(nn.Module):
    def __init__(self, s, k, ic):
        super(Shuffle, self).__init__()
        self.shuffle = nn.PixelShuffle(s)

    def forward(self, inputs):
        return self.shuffle(inputs)

def main():
    sz = 4

    h = 3
    w = 3
    k = 3
    ic = 8
    s = 2

    input = np.arange(0, ic * h * w, dtype=np.float32).reshape(ic, h, w)
    input = input[np.newaxis]
    torch_input = torch.from_numpy(input)

    shuffle_model = Shuffle(s, k, ic)
    shuffle_out = shuffle_model(torch_input).detach().numpy()
    print('Shuffle out:', shuffle_out.shape)
    print(shuffle_out)

    input = np.transpose(input, (0, 2, 3, 1))

    keras_input = keras.layers.Input(shape=input.shape[1:])
    keras_d2s = keras.layers.Lambda(lambda x: tf.nn.depth_to_space(x, s))(keras_input)
    keras_model = keras.Model(keras_input, keras_d2s)
    keras_out = keras_model.predict(input)
    
    keras_out = np.transpose(keras_out, (0, 3, 1, 2))
    
    print('Keras out:', keras_out.shape)
    print(keras_out)
    equal = np.allclose(shuffle_out, keras_out)
    print('Equal?', equal)

if __name__ == '__main__':
    main()

型