add support for freeze training

[youyuxiansen]

This tutorial is for freeze training

For what

There may be a need for some peoples to train a model with some structure freezer.

Modify

1. Modify the --freeze param in train.py，to support more flexible frozen layers selection methods.
2. Add a frozen training yaml file that support more flexible frozen training plan definitions.

How to use

Customize the model training process by defining a yaml file.

Demo

freeze_train.py --weights yolov5s.pt \
--data data/coco128.yaml --cfg models/yolov5s.yaml --batch-size -1 \
--device 0,1,2,3 --hyp data/hyps/hyp.finetune.yaml --cache \
--freeze-plan freeze_plans/freeze_exp.yaml

TODO

• Support selecting part of the data in any freeze training step to train the model.

Implementation

Currently the training process is independent from train.py, because it needs a little bit more time to put it into train.py, and I am not sure if this feature is important for this repository for now. If necessary, I can merge the changes into train.py later, and I will maintain the bugs of this feature in time.

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Implementation of layer freezing mechanism during training in YOLOv5.

📊 Key Changes

• Added freeze_train.py, a new training script to handle freeze training.
• Introduced freeze_plans folder with example YAML files for defining freeze strategies.
• Modified train.py to support the freeze options through --freeze and --freeze-type CLI arguments.
• A README file in freeze_plans detailing how to use freeze training has been added.

🎯 Purpose & Impact

• 🎓 Purpose: Enables training YOLOv5 models with certain layers frozen, enhancing fine-tuning capabilities for transfer learning and potentially improving model performance on specific tasks.
• 👩‍💻 Impact: Allows users more control over the training process, potentially speeding up training and increasing model stability.
• 🦾 To Users: Anyone interested in advanced model training techniques can now experiment with different layer freezing strategies.

[glenn-jocher]

@youyuxiansen you can freeze any layers with the --freeze argument, i.e. to freeze the backbone (first 10 layers):

python train.py --weights yolov5s.pt --freeze 10

See argparser for details:

yolov5/train.py

Line 472 in 628817d

parser.add_argument('--freeze', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24')

[youyuxiansen]

Hi @glenn-jocher , looks like you haven’t read my commitment. I Submit this because the freeze cannot meet some needs. For example, what if I want to freeze the 2,3,4,5 layers or the 1,3,5,7 layer. And if I want to train multiple times with different frozen layers each time. It is so inconvenient and not directly. But with my commitment, this all can be easy to do.

[glenn-jocher]

@youyuxiansen yes the current freeze argument only supports and 'up to layer' value. I think for most customizations a user may want to modify the freeze logic directly here:

yolov5/train.py

Lines 126 to 133 in c1249a4

	# Freeze
	freeze = [f'model.{x}.' for x in range(freeze)] # layers to freeze
	for k, v in model.named_parameters():
	v.requires_grad = True # train all layers
	if any(x in k for x in freeze):
	LOGGER.info(f'freezing {k}')
	v.requires_grad = False

It is true that I think TF supports a wider variety of freezing strategies.

[youyuxiansen]

I suggest you try my submission when you have time. @glenn-jocher

[glenn-jocher]

@youyuxiansen yes I've browsed the code. I understand it adds features. The main issue is the scope and complexity of additional code that the PR introduces, and the addition of a new file in the YOLOv5 root directory (which itself introduces duplication of existing code/functionality).

For various reasons (maintenance, documentation, simplicity) I would encourage you to find the minimum viable solution with the least amount of code required.

[youyuxiansen]

@glenn-jocher Ok, I understand. At last, can this part of the commit be accepted for a more flexible frozen layer specified?

https://github.com/ultralytics/yolov5/pull/6001/files#diff-ed183d67207df065a11e1289f19d34cc2abbc5448dea952683cfe9728c342b95R127-R137

[glenn-jocher]

@youyuxiansen yes I was going to say we might be able to smartly handle arguments so that a single value would be 'up to layer' and multiple values would simply index the freeze layers, i.e.:

python train.py --freeze 10  # freeze up to 10
python train.py --freeze 5 6 7 8 9 10  # freeze layers 5-10

EDIT: we don't want to add additional arguments (unless absolutely necessary for a major feature) to train.py as it already has too many arguments.

EDIT2: see detect.py --weights as an example of variable length argument:

    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')

# python detect.py --weights yolov5s.pt
# python detect.py --weights yolov5s.pt yolov5m.pt  # ensemble

[youyuxiansen]

@glenn-jocher Oh, I see. I think I might have some time tomorrow to modify this for YOLOV5.