First of all, Tensorflow supports two modes of image input channel first and channel behind, while Pytorch only supports the input mode of channel first:
Tensorflow supports two form of data as the input of its convolutional layer, channels first and channels last. In the first format, data are in the form of [N, C, H, W] where N is the batch size, C is the number of channels, H is the height of matrix and W is the width of matrix. In the second format, data are in the form of [N, H, W, C]. In PyTorch only channels first data format is supported.
In Tensorflow, you can set the switching and selection of the two modes. In Pytorch, if the channel is behind, you can only transform the input. You can refer to: https://blog.csdn.net/qq_40877238/article /details/106024442