Torchvision ToTensor, don't change channel order

himat · May 19, 2020, 7:13pm

When you use transforms.ToTensor(), by default it changes the input arrays from HWC to CHW order.

Is it possible to have ToTensor not do this if my input data are already in the ordering I want?

What can I replace this code with to make it keep the ordering I want and still convert to a torch tensor?

transforms.Compose([transforms.ToTensor()])

Nikronic · May 19, 2020, 9:38pm

Hi,

As there is no parameter in transforms.ToTensor() to control this behavior, the easiest and probably efficient way to solve this issue is to do what PyTorch have in its source code.

github.com

pytorch/vision/blob/e6b4078ec73c2cf4fd4432e19c782db58719fb99/torchvision/transforms/functional.py#L78


        img = torch.from_numpy(np.array(pic, np.int16, copy=False))
    elif pic.mode == 'F':
        img = torch.from_numpy(np.array(pic, np.float32, copy=False))
    elif pic.mode == '1':
        img = 255 * torch.from_numpy(np.array(pic, np.uint8, copy=False))
    else:
        img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))

    img = img.view(pic.size[1], pic.size[0], len(pic.getbands()))
    # put it from HWC to CHW format
    img = img.permute((2, 0, 1)).contiguous()
    if isinstance(img, torch.ByteTensor):
        return img.float().div(255)
    else:
        return img


def pil_to_tensor(pic):
    """Convert a ``PIL Image`` to a tensor of the same type.

    See ``AsTensor`` for more details.

As you can see, it permutes the channels, So, same code can be used to reverse the order.

The reason that PyTorch is doing this is that, PyTorch uses CHW convention by default.

bests,
Nik

rayryeng · March 16, 2021, 5:07pm

Thanks for this. Just for the sake of containment if you want to reverse the order done by ToTensor, it would have to be img = img.permute((1, 2, 0)).contiguous(). The (2, 0, 1) is the forward transformation to go from HWC to CHW. To undo the operation, the tuple needs to be (1, 2, 0) instead.

Thanks again!