I want to translate a code from pytorch which uses torch.nn.functional.unfold to tensorflow2. I saw in How to replicate PyTorch's nn.functional.unfold function in Tensorflow? and Pytorch "Unfold" equivalent in Tensorflow that i need to use tf.image.extract_patches() function.
I have:
image = np.random.rand(2,3,32,32)
torch_image = tensor(image)
torch_x = torch.nn.functional.unfold(torch_image, (3,3), dilation=1, padding=0, stride=1)
print(torch_x.shape)
tf_image = tf.convert_to_tensor(image)
tf_image = tf.transpose(tf_image, [0, 2, 3, 1])
tf_x = tf.image.extract_patches(tf_image, sizes=[1,3,3,1], strides=[1,1,1,1], rates=[1,1,1,1], padding="VALID")
print(tf_x.shape)
This code gives me an output torch_x with a shape of (2,27,900) and an output tf_x with a shape of (2,30,30,27). I realize a small test:
a = sorted(list(torch_x.numpy().flatten()))
b = sorted(list(tf_x.numpy().flatten()))
print(set([i-j for i,j in zip(a,b)]))
It results than all the values of tf_x are in torch_x. But, i dont know how to reshape tf_x to be equal to torch_x. I tried :
final_tf_x = tf.transpose(tf_x, [0, 3, 1, 2])
final_tf_x = tf.reshape(final_tf_x, [final_tf_x.shape[0], final_tf_x.shape[1], -1])
print(final_tf_x.shape)
print(np.abs(torch_x.numpy()-final_tf_x.numpy())<1e-8)
It gives me a tensor of the same shape as torch_x but the 2 tensors are not equal elementwise. Can someone explain me how to do this last step?