manually freeze the model before training mask head?

facebookresearch / detr

End-to-End Object Detection with Transformers

Apache License 2.0

13.36k stars 2.41k forks source link

manually freeze the model before training mask head? #413

Open eaglep91 opened 3 years ago

eaglep91 commented 3 years ago

I have trained the box model successfully with 2 object classes, image size is 350x350. There are around 1000 images in training. However when I train the mask head, I'm facing a "CUDA Memory" issue which asks for more memory. I'm wondering if I need to manually freeze the checkpoint before starting the mask head training? If so, is there any proper way of doing that? Thanks!

NielsRogge commented 3 years ago

Hi,

This is done automatically if you specify the --masks argument. You can see this here.

If you face memory issue, try to reduce the batch size.