Currently the reset_classifier function does not work, as there is no self.out_dim in the model.
I've updated reset_classifier to work by using the current head in_features to determine the in channels for the new head.
I also removed global_pool param, which was dead weight and was not being used.
I added a force option in case you are not changing num_classes but still want to reset (defaults to off).
Finally, added a to_gpu flag b/c in building a model that is on cpu atm, you don't want to force only the head to gpu if the entire model will be moved later, which was how the initial version was setup.
ghost
commented
3 years ago
Currently the reset_classifier function does not work, as there is no self.out_dim in the model.
I've updated reset_classifier to work by using the current head in_features to determine the in channels for the new head.
I also removed global_pool param, which was dead weight and was not being used. I added a force option in case you are not changing num_classes but still want to reset (defaults to off). Finally, added a to_gpu flag b/c in building a model that is on cpu atm, you don't want to force only the head to gpu if the entire model will be moved later, which was how the initial version was setup.