search

hhk7734 / tensorflow-yolov4

YOLOv4 Implemented in Tensorflow 2.
MIT License
136 stars 75 forks source link

Transfer learning from COCO with different classes #28

Closed raz-SX closed 3 years ago

raz-SX commented 3 years ago

Hi @hhk7734 thank you for your amazing work!

I'm trying to load the yolov4-tiny-relu.weights into a model with three classes. When I'm doing so, I'm getting a ValueError("Model and weights file do not match.") from the load_weights function.

To my understanding, the number of classes affects the PANet and the YoloHead, but on the weights.py file the distinction between the two in the way the weights are loaded are the layers' indexes.

Could you orient me on how to accommodate loading and saving weights for different class numbers?

Thanks Raz

hhk7734 commented 3 years ago

As you said, the number of classes affects some layers. It changes the size of the feature map. A TF model has info of layers or sub-model. A layer has a unique index or name. So, we can find a layer using an index or name even if the number of classes is changed.

For example, Convolution 17 is in PANetTiny. The name is yolo_conv2d_17, and The filter is 3 * (num_classes + 5).

https://github.com/hhk7734/tensorflow-yolov4/blob/0e73c11ebdfc22f6776b49f6d4dd53415592b585/py_src/yolov4/model/neck.py#L321-L326

The model structure is shown below.

weights.py can find layer info using layer name or index.

In yolo_conv2d_load_weights,

https://github.com/hhk7734/tensorflow-yolov4/blob/0e73c11ebdfc22f6776b49f6d4dd53415592b585/py_src/yolov4/tf/weights.py#L62 https://github.com/hhk7734/tensorflow-yolov4/blob/0e73c11ebdfc22f6776b49f6d4dd53415592b585/py_src/yolov4/tf/weights.py#L84-L87

I don't know if there is a way to extract only the desired classes from the interacting CNN.

raz-SX commented 3 years ago

Thanks!