Pytorch-yolov3: Transfer Learning YOLO

[monteksingh]

I am trying to set the 'requires_grad' property of the specific layers to FALSE to perform Transfer Learning.

But when I introduce the code below in Training.py-

for key, values in model.state_dict().items():
        print(model.state_dict()[key].requires_grad,"Key-",key)

I get the below Output-

False Key- module_list.0.conv_0.weight
False Key- module_list.0.batch_norm_0.weight

i.e 'False' for all the layers even after-> model.train()

I tried to change using ->
Source - https://github.com/ultralytics/yolov3/wiki/Example:-Transfer-Learning

for i, (name, p) in enumerate(model.named_parameters()):
    if (p.shape[0] == 69):
        p.requires_grad = True
        print(p.requires_grad,"pp",name)
    else:
        p.requires_grad = False

69 because I have (4+1+18classes) * 3anchors
Changes are shown in-> print(p.requires_grad,"pp",name) But even after this when I print the 'requires_grad' for individual Layers the output remains FALSE.

for key, values in model.state_dict().items():
        print(model.state_dict()[key].requires_grad,"Key-",key)

Please explain where am I going wrong.......

[selimarikan]

No need to use another repo. Here is how I am doing the transfer learning:

I have this function that prepares the model. It freezes all the layers first up to the provided index.
Then I just create new layers to replace existing ones with correct parameters (class count, attribute count, etc.).
This is a sample for YOLOv3 Tiny but you can easily adapt it for normal YOLO as well.

Another small thing is to visualize the network to find out which layers you want to freeze. I am using netron tool for that. Just load the .CFG file and you will see the network model.

def prepare_pretrained(model, class_count, freeze_index):
    # Modify the network, first freeze all layers
    for name, param in model.named_parameters():
        # module_list.21.batch_norm_21.bias
        layer_id = int(name.split('.')[-3])

        # Freeze layers before 6
        if layer_id < freeze_index:
            param.requires_grad = False

    # For YOLOv3 Tiny, edit these:
    # 23-> yolo_23 -> Last YOLO -> Set classes
    # 22-> conv_22 -> Last Conv -> Set output
    # 16-> yolo_16 -> Prev YOLO ->
    # 15-> conv_15 -> Prev Conv

    # Bounding box attribute count calculation
    # This value is used both in YOLO and last conv layers
    # [4 box coordinates + 1 object confidence + 80 class conf] -> 255 = 85 * 3
    # [4 box coordinates + 1 object confidence + 1 class conf]  -> 18  =  6 * 3
    bb_attrib_count = (5 + class_count) * 3

    # 1- Set last YOLO layer output size
    yolo_23 = model.module_list[23][0]
    new_yolo_23 = YOLOLayer(yolo_23.anchors, class_count, yolo_23.image_dim)
    model.module_list[23][0] = new_yolo_23

    # 2- Set the out filters of conv before last YOLO
    conv_22 = model.module_list[22][0]
    new_conv_22 = torch.nn.Conv2d(
            in_channels=conv_22.in_channels,
                                  out_channels=bb_attrib_count,
                                  kernel_size=conv_22.kernel_size,
                                  stride=conv_22.stride,
                                  padding=conv_22.padding)
    model.module_list[22][0] = new_conv_22
    ...
    ...

[Serhiy-Shekhovtsov]

Thanks for the example @selimarikan, it's very helpful. I just wanted to add my two cents here, image_dim property is now called img_dim. So, you can create new YOLO layer like this:
YOLOLayer(yolo_23.anchors, class_count, yolo_23.img_dim)
Also, probably you will want to move new layers' parameters to a target device:

model.module_list[23][0] = new_yolo_23.to(device)
...
model.module_list[22][0] = new_conv_22.to(device)

[Flova]

I closed this issue due to inactivity. Feel free to reopen for further discussion.