xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Move loss to task heads (#2825)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
This commit is contained in:
Ayush Chaurasia
2023-05-25 16:07:54 +05:30
committed by GitHub
parent f23a03596d
commit d19c5b6ce8
9 changed files with 335 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -325,8 +325,7 @@ class BaseTrainer:
# Forward
with torch.cuda.amp.autocast(self.amp):
batch = self.preprocess_batch(batch)
preds = self.model(batch['img'])
self.loss, self.loss_items = self.criterion(preds, batch)
self.loss, self.loss_items = de_parallel(self.model).loss(batch)
if RANK != -1:
self.loss *= world_size
self.tloss = (self.tloss * i + self.loss_items) / (i + 1) if self.tloss is not None \
@ -496,12 +495,6 @@ class BaseTrainer:
"""Build dataset"""
raise NotImplementedError('build_dataset function not implemented in trainer')
def criterion(self, preds, batch):
"""
Returns loss and individual loss items as Tensor.
"""
raise NotImplementedError('criterion function not implemented in trainer')
def label_loss_items(self, loss_items=None, prefix='train'):
"""
Returns a loss dict with labelled training loss items tensor

3
ultralytics/yolo/engine/validator.py
View File

@ -162,7 +162,8 @@ class BaseValidator:
# Loss
with dt[2]:
if self.training:
self.loss += trainer.criterion(preds, batch)[1]
loss_items = model.loss(batch, preds)
self.loss += loss_items[1]
# Postprocess
with dt[3]: