Как правильно использовать функцию просмотра Pytorch?

Question

У меня есть тензор размера x = {4,2, C, H, W}. Мне нужно, чтобы оно изменилось на y = {8, C, H, W}, но я хочу убедиться, что изображения хранятся в правильном порядке, например, изображение в x [1,0,:,:,:] должен быть равен y [2, C, H, W]. Я знаю, что могу использовать функцию просмотра для этого, но я не уверен, как правильно ее использовать.

В настоящее время я делаю это как таковое

feat_imgs_all = feat_imgs_all.view(
    rgb.shape[0], rgb.shape[1], feat_imgs_all.shape[1], 
    feat_imgs_all.shape[2], feat_imgs_all.shape[3])

Это кажется действительно хакерским, есть ли Как я могу просто кормить первые две фигуры, а Pytorch вычисляет остальные?

Answer 1

Вы можете сделать это легко, используя аргументы flatten и end_dim, см. documentation:

import torch

a = torch.randn(4, 2, 32, 64, 64)
flattened = a.flatten(end_dim=1)

torch.all(flattened[2, ...] == a[1, 0, ...]) # True

view также можно использовать, как показано ниже, хотя это не слишком читабельно и не слишком приятно:

import torch

a = torch.randn(4, 2, 32, 64, 64)
flattened = a.view(-1, *a.shape[2:])

torch.all(flattened[2, ...] == a[1, 0, ...]) # True as well