HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Fix some cuda training issues of segmentation (#46)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Laughing 2022-11-17 06:44:02 -06:00 • committed by GitHub
parent db1031a1a9
commit 47f1cb3ef4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 38 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/yolo/v8/segment/train.py
View file

@ -23,7 +23,7 @@ class SegmentationTrainer(BaseTrainer):
def get_dataloader(self, dataset_path, batch_size, rank=0):
# TODO: manage splits differently
# calculate stride - check if model is initialized
gs = max(int(self.model.stride.max() if self.model else 0), 32)
gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
loader = build_dataloader(
img_path=dataset_path,
img_size=self.args.img_size,
@ -220,7 +220,7 @@ class SegmentationTrainer(BaseTrainer):
mxyxy = xywh2xyxy(xywhn[i] * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device))
for bi in b.unique():
j = b == bi # matching index
if True:
if self.args.overlap_mask:
mask_gti = torch.where(masks[bi][None] == tidxs[i][j].view(-1, 1, 1), 1.0, 0.0)
else:
mask_gti = masks[tidxs[i]][j]