numpy pytorch で画像データを（100,100,3）から （3,100,100）の構造に

Question

　　　　　　　　　　　4,4,2の配列
[[[0. 1.]
  [0. 1.]
  [0. 1.]
  [0. 1.]]

 [[0. 1.]
  [0. 1.]
  [0. 1.]
  [0. 1.]]

 [[0. 1.]
  [0. 1.]
  [0. 1.]
  [0. 1.]]

 [[0. 1.]
  [0. 1.]
  [0. 1.]
  [0. 1.]]]
```
から
```
                          2,4,4の配列
[[[0. 0. 0. 0.]
  [0. 0. 0. 0.]
  [0. 0. 0. 0.]
  [0. 0. 0. 0.]]

 [[1. 1. 1. 1.]
  [1. 1. 1. 1.]
  [1. 1. 1. 1.]
  [1. 1. 1. 1.]]]
```

こうしたいです
できる人やり方教えてください


試したこと
reshapeを使ってもうまくいきません
他の関数はなんて調べたら出てくるかわかんない

Answer 1

Tensor.permute() でしょうか。
numpy.transpose に相当する軸を並び替える関数です。

torch.Tensor — PyTorch 1.5.0 documentation

import torch

x = [
    [[0, 1], [0, 1], [0, 1], [0, 1]],
    [[0, 1], [0, 1], [0, 1], [0, 1]],
    [[0, 1], [0, 1], [0, 1], [0, 1]],
    [[0, 1], [0, 1], [0, 1], [0, 1]],
]

x = torch.Tensor(x)
print(x.shape)

# axis を入れ替える
# before axis : 0  1  2
#        shape: 4  4  2
#    ↓
# after  axis : 2  0  1
#        shape: 2  4  4
x = x.permute(2, 0, 1)

print(x)
# tensor([[[0., 0., 0., 0.],
#          [0., 0., 0., 0.],
#          [0., 0., 0., 0.],
#          [0., 0., 0., 0.]],

#         [[1., 1., 1., 1.],
#          [1., 1., 1., 1.],
#          [1., 1., 1., 1.],
#          [1., 1., 1., 1.]]])

numpy でやりたいことを実現できる関数名を知っていれば「pytorch numpy の関数名」で検索すると、Pytorch でそれに相当する処理はどうやればいいのか出てくると思います。