Skip to content

Commit

Permalink
Fix Readme and remove unused kernel (#270)
Browse files Browse the repository at this point in the history
* fix reamde and remove unused kernel

* remove unused tests

---------
  • Loading branch information
drisspg authored May 24, 2024
1 parent f8f74c7 commit 90b5e17
Show file tree
Hide file tree
Showing 5 changed files with 16 additions and 225 deletions.
16 changes: 15 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,23 @@ git clone https://github.com/pytorch/ao
cd ao
pip install -r requirements.txt
pip install -r dev-requirements.txt
pip install .
```

There are two options;
-If you plan to be developing the library run:
```Shell
python setup.py develop
```

If you want to install from source run
```Shell
python setup.py install
```

** Note:
Since we are building pytorch c++/cuda extensions by default, running `pip install .` will
not work.

### Quantization

```python
Expand Down
15 changes: 0 additions & 15 deletions test/test_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,21 +30,6 @@ def _create_tensors_with_iou(self, N, iou_thresh):
scores = torch.rand(N)
return boxes, scores

@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
@unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "skipping when torch verion is 2.3 or lower")
def test_nms(self):
iou = 0.2
boxes, scores = self._create_tensors_with_iou(1000, iou)
boxes = boxes.cuda()
scores = scores.cuda()

# smoke test
_ = torchao.ops.nms(boxes, scores, iou)

# comprehensive testing
test_utils = ["test_schema", "test_autograd_registration", "test_faketensor", "test_aot_dispatch_dynamic"]
opcheck(torch.ops.torchao.nms, (boxes, scores, iou), test_utils=test_utils)

def _create_fp6_inputs(self, BS: int, OC: int, IC: int):
# Randomly initialize each bytes. The highest value for randint() is set the the max value of uint32_t.
fp6_weight = torch.randint(4294967295, (OC, IC // 16 * 3)).to(torch.int)
Expand Down
181 changes: 0 additions & 181 deletions torchao/csrc/cuda/nms.cu

This file was deleted.

8 changes: 0 additions & 8 deletions torchao/csrc/nms.cpp

This file was deleted.

21 changes: 1 addition & 20 deletions torchao/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,26 +10,6 @@ def decorator(func):
return torch.library.impl_abstract(f"{name}")(func)
return decorator

def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
"""
See https://pytorch.org/vision/main/generated/torchvision.ops.nms.html
"""
return torch.ops.torchao.nms.default(boxes, scores, iou_threshold)


# Defines the meta kernel / fake kernel / abstract impl
@register_custom_op("torchao::nms")
def _(dets, scores, iou_threshold):
torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")
torch._check(scores.dim() == 1, lambda: f"scores should be a 1d tensor, got {scores.dim()}")
torch._check(
dets.size(0) == scores.size(0),
lambda: f"boxes and scores should have same number of elements in dimension 0, got {dets.size(0)} and {scores.size(0)}",
)
ctx = torch._custom_ops.get_ctx()
num_to_keep = ctx.create_unbacked_symint()
return dets.new_empty(num_to_keep, dtype=torch.long)


def prepack_fp6_weight(fp6_weight: Tensor) -> Tensor:
Expand All @@ -45,6 +25,7 @@ def prepack_fp6_weight(fp6_weight: Tensor) -> Tensor:
return torch.ops.torchao.prepack_fp6_weight.default(fp6_weight)


# Defines the meta kernel / fake kernel / abstract impl
@register_custom_op("torchao::prepack_fp6_weight")
def _(fp6_weight):
torch._check(fp6_weight.dim() == 2, lambda: f"weight should be a 2d tensor, got {fp6_weight.dim()}D")
Expand Down

0 comments on commit 90b5e17

Please sign in to comment.