swap tensor axis in keras

Question

I want to swap tensor axis of image batches from (batch_size, row, col, ch) to (batch_size, ch, row, col).

in numpy, this can be done with

X_batch = np.moveaxis( X_batch, 3, 1)

How would I do that in Keras?

Answer 1

You can use K.permute_dimensions() which is exactly similar to np.transpose().

Example:

import numpy as np 
from keras import backend as K 

A = np.random.random((1000,32,64,3))
# B = np.moveaxis( A, 3, 1)
C = np.transpose( A, (0,3,1,2))

print A.shape
print C.shape

A_t = K.variable(A)
C_t = K.permute_dimensions(A_t, (0,3,1,2))

print K.eval(A_t).shape
print K.eval(C_t).shape