osmr / imgclsmob

Sandbox for training deep learning networks

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

How to user pretrained weights with classes != 1000 ?

veryviolet opened this issue · comments

Suppose I want to get EfficientNet B0 with classes=2 and pretrained weights (excluding classifying layer):

from tf2cv.model_provider import get_model

net = get_model(name="efficientnet_b0", pretrained=True, classes=2)

Code above failes with:

ValueError: Shapes (1280, 2) and (1280, 1000) are incompatible

Obviously, weights are loaded for classification layer also, and that leads to an error. How can I avoid loading weights for a classification layer?

There are at least two ways:

  1. After your unsuccessful attempt to load the pre-trained weights, the corresponding weight file was cached in .tensorflow/models. You can load these weighs by using method net.load_weights() with appropriate parameters.
  2. You can load the feature extractor efficientnet_b0(pretrained=True).features and use it in your own model with any classifier.