xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests 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
2 changed files with 17 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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)