Skip to content

Commit 90b5e17

Browse files
authored
Fix Readme and remove unused kernel (#270)
* fix reamde and remove unused kernel * remove unused tests ---------
1 parent f8f74c7 commit 90b5e17

File tree

5 files changed

+16
-225
lines changed

5 files changed

+16
-225
lines changed

README.md

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,23 @@ git clone https://github.com/pytorch/ao
2929
cd ao
3030
pip install -r requirements.txt
3131
pip install -r dev-requirements.txt
32-
pip install .
3332
```
3433

34+
There are two options;
35+
-If you plan to be developing the library run:
36+
```Shell
37+
python setup.py develop
38+
```
39+
40+
If you want to install from source run
41+
```Shell
42+
python setup.py install
43+
```
44+
45+
** Note:
46+
Since we are building pytorch c++/cuda extensions by default, running `pip install .` will
47+
not work.
48+
3549
### Quantization
3650

3751
```python

test/test_ops.py

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -30,21 +30,6 @@ def _create_tensors_with_iou(self, N, iou_thresh):
3030
scores = torch.rand(N)
3131
return boxes, scores
3232

33-
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
34-
@unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "skipping when torch verion is 2.3 or lower")
35-
def test_nms(self):
36-
iou = 0.2
37-
boxes, scores = self._create_tensors_with_iou(1000, iou)
38-
boxes = boxes.cuda()
39-
scores = scores.cuda()
40-
41-
# smoke test
42-
_ = torchao.ops.nms(boxes, scores, iou)
43-
44-
# comprehensive testing
45-
test_utils = ["test_schema", "test_autograd_registration", "test_faketensor", "test_aot_dispatch_dynamic"]
46-
opcheck(torch.ops.torchao.nms, (boxes, scores, iou), test_utils=test_utils)
47-
4833
def _create_fp6_inputs(self, BS: int, OC: int, IC: int):
4934
# Randomly initialize each bytes. The highest value for randint() is set the the max value of uint32_t.
5035
fp6_weight = torch.randint(4294967295, (OC, IC // 16 * 3)).to(torch.int)

torchao/csrc/cuda/nms.cu

Lines changed: 0 additions & 181 deletions
This file was deleted.

torchao/csrc/nms.cpp

Lines changed: 0 additions & 8 deletions
This file was deleted.

torchao/ops.py

Lines changed: 1 addition & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -10,26 +10,6 @@ def decorator(func):
1010
return torch.library.impl_abstract(f"{name}")(func)
1111
return decorator
1212

13-
def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
14-
"""
15-
See https://pytorch.org/vision/main/generated/torchvision.ops.nms.html
16-
"""
17-
return torch.ops.torchao.nms.default(boxes, scores, iou_threshold)
18-
19-
20-
# Defines the meta kernel / fake kernel / abstract impl
21-
@register_custom_op("torchao::nms")
22-
def _(dets, scores, iou_threshold):
23-
torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
24-
torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")
25-
torch._check(scores.dim() == 1, lambda: f"scores should be a 1d tensor, got {scores.dim()}")
26-
torch._check(
27-
dets.size(0) == scores.size(0),
28-
lambda: f"boxes and scores should have same number of elements in dimension 0, got {dets.size(0)} and {scores.size(0)}",
29-
)
30-
ctx = torch._custom_ops.get_ctx()
31-
num_to_keep = ctx.create_unbacked_symint()
32-
return dets.new_empty(num_to_keep, dtype=torch.long)
3313

3414

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

4727

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

0 commit comments

Comments
 (0)