efficientnetv2 export issue trained for custom dataset

[bnascimento]

Hi everyone, I trained / fine tuned an efficientnetv2-s from a custom dataset, without a problem. I changed the datasets.py with the necessary changes (my model was for example 278 classes) and I am now trying to export to it saved_model using the existing code. So, when running

python3 infer.py --mode tf2bm --model_name efficientnetv2-s --model_dir ../train_efficientnetv2-s/ --export_dir ../train_efficientnetv2-s/saved_model

I am getting the following error

Traceback (most recent call last): File "infer.py", line 203, in <module> app.run(main) File "/usr/local/lib/python3.8/dist-packages/absl/app.py", line 312, in run _run_main(main, args) File "/usr/local/lib/python3.8/dist-packages/absl/app.py", line 258, in _run_main sys.exit(main(argv)) File "infer.py", line 193, in main tf2_benchmark() File "infer.py", line 91, in tf2_benchmark model = build_tf2_model() File "infer.py", line 55, in build_tf2_model model = effnetv2_model.get_model( File "/opt/automl/efficientnetv2/effnetv2_model.py", line 787, in get_model net.load_weights(pretrained_ckpt) File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler raise e.with_traceback(filtered_tb) from None File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/training/saving/saveable_object_util.py", line 134, in restore raise ValueError( ValueError: Received incompatible tensor with shape (278,) when attempting to restore variable with shape (1000,) and name _fc/bias/.ATTRIBUTES/VARIABLE_VALUE.

So the question is how can I export a checkpoint with a customized number of outputs? Instead of 1000 ImageNet outputs, 278 for examples. I tried to use the FLAGS.dataset_cfg, but looking at the code that does not seem to be used on infer.py Is there any --hparam_str param we can overwrite?

Thanks all

[bnascimento]

was able to solve it with --hparam_str="num_classes=278"