From 2db35afad51f5ab0e327e18d3bf570b0aad968fd Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 26 Aug 2023 17:27:18 +0200 Subject: [PATCH] Simplify `Results()` class (#4579) --- docker/Dockerfile | 5 +++- docs/quickstart.md | 3 ++- tests/test_python.py | 15 +++++------- ultralytics/engine/results.py | 46 ++++++++++++++--------------------- 4 files changed, 30 insertions(+), 39 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 81b9f99568..24e09cd179 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -52,9 +52,12 @@ ENV MKL_THREADING_LAYER=GNU # Build and Push # t=ultralytics/ultralytics:latest && sudo docker build -f docker/Dockerfile -t $t . && sudo docker push $t -# Pull and Run +# Pull and Run with access to all GPUs # t=ultralytics/ultralytics:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t +# Pull and Run with access to GPUs 2 and 3 (inside container CUDA devices will appear as 0 and 1) +# t=ultralytics/ultralytics:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus '"device=2,3"' $t + # Pull and Run with local directory access # t=ultralytics/ultralytics:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/datasets:/usr/src/datasets $t diff --git a/docs/quickstart.md b/docs/quickstart.md index 4d5f6eb8ea..2a6bd5d79d 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -82,7 +82,8 @@ Ultralytics provides various installation methods including pip, conda, and Dock sudo docker pull $t # Run the ultralytics image in a container with GPU support - sudo docker run -it --ipc=host --gpus all $t + sudo docker run -it --ipc=host --gpus all $t # all GPUs + sudo docker run -it --ipc=host --gpus '"device=2,3"' $t # specify GPUs ``` The above command initializes a Docker container with the latest `ultralytics` image. The `-it` flag assigns a pseudo-TTY and maintains stdin open, enabling you to interact with the container. The `--ipc=host` flag sets the IPC (Inter-Process Communication) namespace to the host, which is essential for sharing memory between processes. The `--gpus all` flag enables access to all available GPUs inside the container, which is crucial for tasks that require GPU computation. diff --git a/tests/test_python.py b/tests/test_python.py index eb418cd7b5..159ca8b896 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -217,7 +217,7 @@ def test_all_model_yamls(): for m in (ROOT / 'cfg' / 'models').rglob('*.yaml'): if 'rtdetr' in m.name: if TORCH_1_9: # torch<=1.8 issue - TypeError: __init__() got an unexpected keyword argument 'batch_first' - RTDETR(m.name)(SOURCE, imgsz=640) + RTDETR(m.name)(SOURCE, imgsz=640) # must be 640 else: YOLO(m.name) @@ -225,8 +225,8 @@ def test_all_model_yamls(): def test_workflow(): model = YOLO(MODEL) model.train(data='coco8.yaml', epochs=1, imgsz=32) - model.val() - model.predict(SOURCE) + model.val(imgsz=32) + model.predict(SOURCE, imgsz=32) model.export(format='onnx') # export a model to ONNX format @@ -243,7 +243,7 @@ def test_predict_callback_and_setup(): dataset = load_inference_source(source=SOURCE) bs = dataset.bs # noqa access predictor properties - results = model.predict(dataset, stream=True) # source already setup + results = model.predict(dataset, stream=True, imgsz=160) # source already setup for r, im0, bs in results: print('test_callback', im0.shape) print('test_callback', bs) @@ -254,7 +254,7 @@ def test_predict_callback_and_setup(): def test_results(): for m in 'yolov8n-pose.pt', 'yolov8n-seg.pt', 'yolov8n.pt', 'yolov8n-cls.pt': model = YOLO(m) - results = model([SOURCE, SOURCE]) + results = model([SOURCE, SOURCE], imgsz=160) for r in results: r = r.cpu().numpy() r = r.to(device='cpu', dtype=torch.float32) @@ -263,10 +263,7 @@ def test_results(): r.tojson(normalize=True) r.plot(pil=True) r.plot(conf=True, boxes=True) - print(r) - print(r.path) - for k in r.keys: - print(getattr(r, k)) + print(r, len(r), r.path) @pytest.mark.skipif(not ONLINE, reason='environment is offline') diff --git a/ultralytics/engine/results.py b/ultralytics/engine/results.py index 55ae280fae..d636395f03 100644 --- a/ultralytics/engine/results.py +++ b/ultralytics/engine/results.py @@ -101,19 +101,18 @@ class Results(SimpleClass): self.names = names self.path = path self.save_dir = None - self._keys = ('boxes', 'masks', 'probs', 'keypoints') + self._keys = 'boxes', 'masks', 'probs', 'keypoints' def __getitem__(self, idx): """Return a Results object for the specified index.""" - r = self.new() - for k in self.keys: - setattr(r, k, getattr(self, k)[idx]) - return r + return self._apply('__getitem__', idx) def __len__(self): """Return the number of detections in the Results object.""" - for k in self.keys: - return len(getattr(self, k)) + for k in self._keys: + v = getattr(self, k) + if v is not None: + return len(v) def update(self, boxes=None, masks=None, probs=None): """Update the boxes, masks, and probs attributes of the Results object.""" @@ -125,43 +124,34 @@ class Results(SimpleClass): if probs is not None: self.probs = probs - def cpu(self): - """Return a copy of the Results object with all tensors on CPU memory.""" + def _apply(self, fn, *args, **kwargs): r = self.new() - for k in self.keys: - setattr(r, k, getattr(self, k).cpu()) + for k in self._keys: + v = getattr(self, k) + if v is not None: + setattr(r, k, getattr(v, fn)(*args, **kwargs)) return r + def cpu(self): + """Return a copy of the Results object with all tensors on CPU memory.""" + return self._apply('cpu') + def numpy(self): """Return a copy of the Results object with all tensors as numpy arrays.""" - r = self.new() - for k in self.keys: - setattr(r, k, getattr(self, k).numpy()) - return r + return self._apply('numpy') def cuda(self): """Return a copy of the Results object with all tensors on GPU memory.""" - r = self.new() - for k in self.keys: - setattr(r, k, getattr(self, k).cuda()) - return r + return self._apply('cuda') def to(self, *args, **kwargs): """Return a copy of the Results object with tensors on the specified device and dtype.""" - r = self.new() - for k in self.keys: - setattr(r, k, getattr(self, k).to(*args, **kwargs)) - return r + return self._apply('to', *args, **kwargs) def new(self): """Return a new Results object with the same image, path, and names.""" return Results(orig_img=self.orig_img, path=self.path, names=self.names) - @property - def keys(self): - """Return a list of non-empty attribute names.""" - return [k for k in self._keys if getattr(self, k) is not None] - def plot( self, conf=True,