If you're looking at the fastest option, go for .transpose(...)
. It's even faster than np.einsum
.
img = np.random.random((1000, 1000, 3))
img.shape
# (1000, 1000, 3)
%timeit img.transpose(2, 0, 1)
# 385 ns ± 1.11 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
%timeit np.rollaxis(img, -1, 0)
# 2.7 µs ± 50.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit np.einsum('ijk->kij', img)
# 2.75 µs ± 31.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit np.moveaxis(img, -1, 0)
# 7.26 µs ± 57.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
np.allclose(img.transpose(2, 0, 1), np.einsum('ijk->kij', img))
# True
np.allclose(img.transpose(2, 0, 1), np.moveaxis(img, -1, 0))
# True
np.allclose(img.transpose(2, 0, 1), np.rollaxis(img,-1, 0))
# True
img.transpose(2,0,1)
orimg.transpose(2,1,0)
. – Donnellytorchvision
, you can also usetorchvision.transforms.ToTensor()(img)
. It will do that automatically for you. – Backswepttorchvision.transforms.ToTensor()(img)
scales your images, and sometimes you don't want that. – Laminar