google / automl

Google Brain AutoML
Apache License 2.0
6.18k stars 1.45k forks source link

export saved_model in NCHW format #1162

Closed bnascimento closed 1 year ago

bnascimento commented 1 year ago

Hi everyone, bumped into an issue. I need to export a efficientnetv2 saved_model to the NCHW format. Is this possible? When I trained my custom model with

--hparam_str="model.data_format=channels_first"

I get the following error

ValueError: The channel dimension of the inputs should be defined. The input_shape received is (None, None, None, 3), where axis -3 (0-based) is the channel dimension, which found to be 'None'.

mingxingtan commented 1 year ago

If you use NCHW, then your inputs should be (None, 3, None, None), where 3 is the input channel size.

bnascimento commented 1 year ago

Hi @mingxingtan , thanks for your reply. Which parameter do I have to overwrite to correctly setup the input format? Either on train time (using main_tf2.py) or export time (using infer.py), is not important, I just need a way to do it.

--hparam_str="model.data_format=channels_first" alone does not seem to be sufficient.

mingxingtan commented 1 year ago

Add support for channels_first. Note, it is not well tested and we generally only use channels_last.