YOLOv5 Apple Metal Performance Shader (MPS) support (#7878)
* Apple Metal Performance Shader (MPS) device support Following https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/ Should work with Apple M1 devices with PyTorch nightly installed with command `--device mps`. Usage examples: ```bash python train.py --device mps python detect.py --device mps python val.py --device mps ``` * Update device strategy to fix MPS issue
This commit is contained in:
1 parent
27911dc824
commit
c215878f11
5 files changed
+11
-8
No files matched your search
@@ -486,7 +486,7 @@ def run(
|
||||
if half:
|
||||
assert device.type != 'cpu' or coreml or xml, '--half only compatible with GPU export, i.e. use --device 0'
|
||||
assert not dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both'
|
||||
model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
|
||||
model = attempt_load(weights, device=device, inplace=True, fuse=True) # load FP32 model
|
||||
nc, names = model.nc, model.names # number of classes, class names
|
||||
|
||||
# Checks
|
||||
|
||||
Reference in new issue
Block a user