def unsqueeze(input, upscale_factor=2): return torch.pixel_shuffle(input, upscale_factor)
def test_pixel_shuffle(self): x = torch.randn(2, 8, 3, 4).float() self.assertONNX(lambda x: torch.pixel_shuffle(x, upscale_factor=2), x, opset_version=11)
def equivariant_unsqueeze(input, upscale_factor=2): in_type = input.type output = torch.pixel_shuffle(input.tensor, upscale_factor) return enn.GeometricTensor(output, in_type)
def forward(self, x): return torch.pixel_shuffle(x, upscale_factor=2)
import torch a = torch.randn(100, 16, 224, 224) b = torch.pixel_shuffle(a, 2) print(b.shape) # torch.Size([100, 4, 448, 448]) c = torch.nn.PixelShuffle(2) d = c(a) print(d.shape) # torch.Size([100, 4, 448, 448])