xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 13:34:23 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Allow setting model attributes before training (#45)

Browse Source
This commit is contained in:
Ayush Chaurasia 2022-11-17 10:35:53 +05:30 committed by GitHub
parent 832ea56eb4
commit db1031a1a9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 17 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
ultralytics/yolo/engine/trainer.py
View File

@ -133,6 +133,7 @@ class BaseTrainer:
"""
Builds dataloaders and optimizer on correct rank process
"""
self.set_model_attributes()
self.optimizer = build_optimizer(model=self.model,
name=self.args.optimizer,
lr=self.args.lr0,
@ -146,19 +147,6 @@ class BaseTrainer:
print("created testloader :", rank)
self.console.info(self.progress_string())
def _set_model_attributes(self):
# TODO: fix and use after self.data_dict is available
'''
head = utils.torch_utils.de_parallel(self.model).model[-1]
self.args.box *= 3 / head.nl # scale to layers
self.args.cls *= head.nc / 80 * 3 / head.nl # scale to classes and layers
self.args.obj *= (self.args.img_size / 640) ** 2 * 3 / nl # scale to image size and layers
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
'''
def _do_train(self, rank, world_size):
if world_size > 1:
self._setup_ddp(rank, world_size)
@ -302,6 +290,12 @@ class BaseTrainer:
if not self.best_fitness or self.best_fitness < self.fitness:
self.best_fitness = self.fitness
def set_model_attributes(self):
"""
To set or update model parameters before training.
"""
pass
def build_targets(self, preds, targets):
pass

10
ultralytics/yolo/v8/segment/train.py
View File

@ -54,6 +54,16 @@ class SegmentationTrainer(BaseTrainer):
model.load(weights)
return model
def set_model_attributes(self):
nl = de_parallel(self.model).model[-1].nl # number of detection layers (to scale hyps)
self.args.box *= 3 / nl # scale to layers
self.args.cls *= self.data["nc"] / 80 * 3 / nl # scale to classes and layers
self.args.obj *= (self.args.img_size / 640) ** 2 * 3 / nl # scale to image size and layers
self.model.nc = self.data["nc"] # attach number of classes to model
self.model.args = self.args # attach hyperparameters to model
# TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
self.model.names = self.data["names"]
def get_validator(self):
return v8.segment.SegmentationValidator(self.test_loader, self.device, logger=self.console)