WongKinYiu / PyTorch_YOLOv4

PyTorch implementation of YOLOv4
1.89k stars 586 forks source link

Use Transfer Learning YoloV4 and yolov4 tiny #382

Open anantgupta129 opened 2 years ago

anantgupta129 commented 2 years ago

How can I use pre trained weights on COCO dataset and finetune them on my dataset for yolov4 and yolov4 tiny? While loading the model only Pytorch weights are used.

Can i add this in https://github.com/WongKinYiu/PyTorch_YOLOv4/blob/master/train.py#L86

    pretrained = weights.endswith('.pt')
    if pretrained:
        with torch_distributed_zero_first(rank):
            attempt_download(weights)  # download if not found locally
        ckpt = torch.load(weights, map_location=device)  # load checkpoint
        model = Darknet(opt.cfg).to(device)  # create
        state_dict = {k: v for k, v in ckpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
        model.load_state_dict(state_dict, strict=False)
        print('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights))  # report
    else:
        model = Darknet(opt.cfg).to(device) # create
         load_darknet_weights(model, 'yolov4-tiny.weights')

@WongKinYiu @glenn-jocher

ingbeeedd commented 2 years ago

Is it work?

parthkvv commented 2 years ago

Is it working?