Convert a tensorflow model to PyTorch

vision
1

Hi,
I need to convert a model with PyTorch codes.
CLASSES = 5
input_size = 224
channels = 3
Here is the tensorflow code:

def build_feature_extractor(inputs):

x = tf.keras.layers.Conv2D(16, kernel_size=3, activation='relu', input_shape=(input_size, input_size, channels))(inputs)
x = tf.keras.layers.AveragePooling2D(2,2)(x)

x = tf.keras.layers.Conv2D(32, kernel_size=3, activation = 'relu')(x)
x = tf.keras.layers.AveragePooling2D(2,2)(x)

x = tf.keras.layers.Conv2D(64, kernel_size=3, activation = 'relu')(x)
x = tf.keras.layers.AveragePooling2D(2,2)(x)

return x

Now I need to convert it to PyTorch.
Conv2d in TensorFlow receive a parameter => input_shape=( *, *, *)
Is it necessary to define input_shape in PyTorch? If the answer is yes, then how can I define it?

Thank you

2

You do not need to specify the input shape to define a Conv2D layer in PyTorch, only the number of input channels, as that is all that is needed to determine the shape of the weights.

See the Conv2D docs here: Conv2d — PyTorch 2.0 documentation which contains information about which arguments are required to define the layer.

3

That means Conv2d in PyTorch is flexible and I don’t need to specify the image size and width.
Thank you eqy.