Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

Avoid FP64 ops for MPS support in train.py (#8511)

Browse Source 
Avoid FP64 ops for MPS support

Resolves https://github.com/ultralytics/yolov5/pull/7878#issuecomment-1177952614
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-07-07 20:36:23 +02:00
1 parent 9d7bc06ae7
commit dd28df98c2
1 file changed
+3 -3
Download Patch File Download Diff File Expand all files Collapse all files
utils/general.py
+3 -3
View File
@@ -644,7 +644,7 @@ def labels_to_class_weights(labels, nc=80):
return torch.Tensor() return torch.Tensor()
labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO
classes = labels[:, 0].astype(np.int) # labels = [class xywh] classes = labels[:, 0].astype(int) # labels = [class xywh]
weights = np.bincount(classes, minlength=nc) # occurrences per class weights = np.bincount(classes, minlength=nc) # occurrences per class
# Prepend gridpoint count (for uCE training) # Prepend gridpoint count (for uCE training)
@@ -654,13 +654,13 @@ def labels_to_class_weights(labels, nc=80):
weights[weights == 0] = 1 # replace empty bins with 1 weights[weights == 0] = 1 # replace empty bins with 1
weights = 1 / weights # number of targets per class weights = 1 / weights # number of targets per class
weights /= weights.sum() # normalize weights /= weights.sum() # normalize
return torch.from_numpy(weights) return torch.from_numpy(weights).float()
def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)): def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
# Produces image weights based on class_weights and image contents # Produces image weights based on class_weights and image contents
# Usage: index = random.choices(range(n), weights=image_weights, k=1) # weighted image sample # Usage: index = random.choices(range(n), weights=image_weights, k=1) # weighted image sample
class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels]) class_counts = np.array([np.bincount(x[:, 0].astype(int), minlength=nc) for x in labels])
return (class_weights.reshape(1, nc) * class_counts).sum(1) return (class_weights.reshape(1, nc) * class_counts).sum(1)