Predictions from SimCLR in the Model Garden

Hi all,

I have been experimenting with the semisupervised learning algorithm SimCLR as implementation in the Model Garden. It is under official.vision.beta.projects.simclr. I am able to successfully pretrain and finetune the model on custom data and the summary metrics look extremely promising. However, I cannot figure out how to load a trained model and generate predictions on new data. I would really appreciate some help.

I think that one specific problem that I’m having is generating or loading the checkpoints generated during training. Let’s say I do a training run. I get logs of the validation metrics and the fully-trained model checkpoint. If I then load the checkpoint and compute performance metrics again on my validation data, I get vastly different numbers. Validation metrics go from very good to no better than chance. It seems to me that either the checkpoint doesn’t store all the weights that the model needs, or that I’m loading the checkpoint wrong. Can anyone comment?

I have tried a number of methods to save the full model to disk once it is finished training. None of them work. I would appreciate some insight here if anyone has any.

  • The model object has its own model.save() method. I believe this is inherited from Keras. This method fails with a complaint about the SimCLR resnet backbone being incompatible somehow with what Keras expects. On Cifar10 for example, I get a message that the model I’m trying to save has an input size of [32,32,6] which is incompatible with the expected shape of [32,32,3].
  • The TensorFlow Model Garden library has its own export method tfm.core.export_base.export(). This also fails, I believe with a complaint that the method wasn’t implemented by SimCLR.
  • I have tried tf.saved_model.save(), which also fails.

Thanks for any and all help.