diff --git a/README.md b/README.md
index 39d690e6..a601ff06 100644
--- a/README.md
+++ b/README.md
@@ -49,7 +49,7 @@ Models from this repo are faster to train (**single card trainable**) and often
| lane detection | ERFNet, ENet | [SAD](https://arxiv.org/abs/1908.00821) ([*Postponed*](https://github.com/voldemortX/pytorch-auto-drive/wiki/Notes)) |
| lane detection | ERFNet | [PRNet](http://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123630698.pdf) (*In progress*) |
| lane detection | ResNets (18, 34, 50, 101), ResNet18-reduced | [LSTR](https://arxiv.org/abs/2011.04233) |
-| lane detection | ResNets (18, 34) | [LaneAtt](https://arxiv.org/abs/2010.12035) |
+| lane detection | ResNets (18, 34) | [LaneATT](https://arxiv.org/abs/2010.12035) |
| lane detection | ResNets (18, 34) | [BézierLaneNet](/configs/lane_detection/bezierlanenet) |
## Model Zoo
diff --git a/docs/DEPLOY.md b/docs/DEPLOY.md
index 317a369a..6cab3879 100644
--- a/docs/DEPLOY.md
+++ b/docs/DEPLOY.md
@@ -24,6 +24,12 @@ In this version, TensorRT may use CUDA runtime >= 11, you might avoid using cond
Or you can incrementally install **Extra Dependencies** through the tutorial.
+## Important Note
+
+Note that we only convert the model `forward()` function,
+post-processing (i.e., `inference()`) is not included. Typical post-processing includes: segmentation map interpolation,
+line NMS for anchor-based lane detection, sigmoid/softmax activations, etc.
+
## PyTorch -> ONNX:
**PyTorch version >= 1.6.0 is recommended for this feature.**
@@ -52,7 +58,7 @@ Same config mechanism and commandline overwrite by `--cfg-options` as in trainin
For detailed instructions and commandline shortcuts available, run:
```
-python tools/profiling.py --help
+python tools/to_onnx.py --help
```
### Currently Unsupported Models:
@@ -89,4 +95,4 @@ You'll then see the saved `ckpt.engine` file and a report on the conversion qual
- SCNN (lane detection)
- Swin backbone (lane detection)
- DCNv2 in BézierLaneNet (lane detection, could support if built custom op from mmcv and directly convert from PyTorch to TensorRT)
-- LaneATT (lane detection)
+- LaneATT (supported if TensorRT >= 8.4.1.5)
diff --git a/docs/MODEL_ZOO.md b/docs/MODEL_ZOO.md
index de2c05a7..6a80215c 100644
--- a/docs/MODEL_ZOO.md
+++ b/docs/MODEL_ZOO.md
@@ -7,7 +7,7 @@
- **level 0**: only small rotation and resize
- **level 1a**: the LSTR augmentations
- **level 1b**: the BézierLaneNet augmentations
- - **level 1c**: the LaneAtt augmentations
+ - **level 1c**: the LaneATT augmentations
| method | backbone | data
augmentation | resolution | mixed precision? | dataset | metric | average | best | training time
*(2080 Ti)* |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
@@ -73,8 +73,8 @@
| LSTR | ResNet18s-2X# | level 0 | 288 x 800 | *no* | CULane | F1 | 36.27 | 39.77 | 28.5h* |
| LSTR | ResNet18s-2X# | level 1a | 288 x 800 | *no* | CULane | F1 | 68.35 | 68.72 | 31.5h* |
| LSTR | ResNet34 | level 1a | 288 x 800 | *no* | CULane | F1 | 72.17 | 72.48 | 45.0h* |
-| LaneAtt | ResNet18 | level 1c | 360 x 640 | *no* | CULane | F1 | 74.71 | 74.87 | 3.6h** |
-| LaneAtt | ResNet34 | level 1c | 360 x 640 | *no* | CULane | F1 | 75.76 | 75.82 | 4.0h** |
+| LaneATT | ResNet18 | level 1c | 360 x 640 | *no* | CULane | F1 | 74.71 | 74.87 | 3.6h** |
+| LaneATT | ResNet34 | level 1c | 360 x 640 | *no* | CULane | F1 | 75.76 | 75.82 | 4.0h** |
| BézierLaneNet | ResNet18 | level 1b | 288 x 800 | *yes* | CULane | F1 | 73.36 | 73.67 | 9.9h |
| BézierLaneNet | ResNet34 | level 1b | 288 x 800 | *yes* | CULane | F1 | 75.30 | 75.57 | 11.0h |
| Baseline | ERFNet | level 0 | 360 x 640 | *yes* | LLAMAS | F1 | 95.94 | 96.13 | 10.9h+ |
@@ -169,8 +169,8 @@
| LSTR | ResNet18s-2X | level 0 | 56.17 | 39.10 | 22.90 | 25.62 | 25.49 | 52.09 | 40.21 | 30.33 | 1690 | 39.77 | [model](https://drive.google.com/file/d/1vdYwM0xDcQLjMAibjmls8hX-IsUe0xcq/view?usp=sharing) \| [shell](../tools/shells/resnet18s_lstr_culane.sh) |
| LSTR | ResNet18s-2X | level 1a | 86.78 | 67.34 | 59.92 | 40.10 | 59.82 | 78.66 | 56.63 | 56.64 | 1166 | 68.72 | [model](https://drive.google.com/file/d/11Tv_nowlWmQtTYQfhGsziDIzb20kPo8o/view?usp=sharing) \| [shell](../tools/shells/resnet18s_lstr-aug_culane.sh) |
| LSTR | ResNet34 | level 1a | 89.73 | 69.77 | 66.72 | 45.32 | 68.16 | 85.03 | 64.34 | 64.13 | 1247 | 72.48 | [model](https://drive.google.com/file/d/1KfmXubuAtUoE9MO8iViMyB_3XhTxWnwH/view?usp=sharing) \| [shell](../tools/shells/resnet34_lstr-aug_culane.sh) |
-| LaneAtt | ResNet18 | level 1c | 90.74 | 72.63 | 69.53 | 47.71 | 70.38 | 86.55 | 65.02 | 65.73 | 1036 | 74.87 | [model](https://drive.google.com/file/d/17Ve7RiqxHK4aEXOt7MJzycj6mvvrieZ1/view?usp=sharing) \| [shell](../tools/shells/resnet18_laneatt_culane.sh) |
-| LaneAtt | ResNet34 | level 1c | 91.36 | 73.72 | 70.71 | 48.40 | 73.69 | 86.86 | 68.95 | 66.00 | 965 | 75.82 | [model](https://drive.google.com/file/d/1AEyNZFAskPg2MKcD3KCaL95JsBeJG_48/view?usp=sharing) \| [shell](../tools/shells/resnet34_laneatt_culane.sh) |
+| LaneATT | ResNet18 | level 1c | 90.74 | 72.63 | 69.53 | 47.71 | 70.38 | 86.55 | 65.02 | 65.73 | 1036 | 74.87 | [model](https://drive.google.com/file/d/17Ve7RiqxHK4aEXOt7MJzycj6mvvrieZ1/view?usp=sharing) \| [shell](../tools/shells/resnet18_laneatt_culane.sh) |
+| LaneATT | ResNet34 | level 1c | 91.36 | 73.72 | 70.71 | 48.40 | 73.69 | 86.86 | 68.95 | 66.00 | 965 | 75.82 | [model](https://drive.google.com/file/d/1AEyNZFAskPg2MKcD3KCaL95JsBeJG_48/view?usp=sharing) \| [shell](../tools/shells/resnet34_laneatt_culane.sh) |
| BézierLaneNet | ResNet18 | level 1b | 90.22 | 71.55 | 68.70 | 45.30 | 70.91 | 84.09 | 62.49 | 58.98 | 996 | 73.67 | [model](https://drive.google.com/file/d/1IpfusHvFeMEGe8wv0fer6KF3pH4X2Tj3/view?usp=sharing) \| [shell](../tools/shells/resnet18_bezierlanenet_culane-aug1b.sh) |
| BézierLaneNet | ResNet34 | level 1b | 91.59 | 73.20 | 69.90 | 48.05 | 76.74 | 87.16 | 69.20 | 62.45 | 888 | 75.57 | [model](https://drive.google.com/file/d/1342FQeDQKRHMo283jW2T1WDgfgsYbR5q/view?usp=sharing) \| [shell](../tools/shells/resnet34_bezierlanenet_culane-aug1b.sh) |
diff --git a/docs/MODEL_ZOO_J.md b/docs/MODEL_ZOO_J.md
index 06ac716a..d32b7489 100644
--- a/docs/MODEL_ZOO_J.md
+++ b/docs/MODEL_ZOO_J.md
@@ -6,7 +6,7 @@ The whole [MODEL_ZOO](./MODEL_ZOO.md) with specifications.
- SCNN
- RESA
- LSTR
-- LaneAtt
+- LaneATT
- [BézierLaneNet](/configs/lane_detection/bezierlanenet)
## Semantic Segmentation
diff --git a/utils/models/lane_detection/laneatt.py b/utils/models/lane_detection/laneatt.py
index afaa1e47..51c4c1c9 100644
--- a/utils/models/lane_detection/laneatt.py
+++ b/utils/models/lane_detection/laneatt.py
@@ -42,7 +42,8 @@ def __init__(self,
anchor_feat_channels=None,
conf_thres=None,
nms_thres=0,
- nms_topk=3000):
+ nms_topk=3000,
+ trace_arg=None):
super().__init__()
self.backbone = MODELS.from_dict(backbone_cfg)
self.backbone_channels = backbone_channels
@@ -61,6 +62,10 @@ def __init__(self,
self.nms_thres = nms_thres
self.nms_topk = nms_topk
+ if trace_arg is not None: # Pre-compute
+ attention_matrix = torch.eye(topk_anchors).repeat(trace_arg['bs'], 1, 1)
+ self.pre_non_diag_inds = torch.nonzero(attention_matrix == 0., as_tuple=False)
+
# generate anchors
self.anchors, self.anchors_cut = self.generate_anchors(lateral_n=72, bottom_n=128)
# Filter masks if `anchors_freq_path` is provided
@@ -209,18 +214,12 @@ def forward(self, x):
attention = softmax(scores).reshape(x.shape[0], len(self.anchors), -1)
attention_matrix = torch.eye(attention.shape[1], device=x.device).repeat(x.shape[0], 1, 1)
if is_tracing():
- # TODO: this also triggers nonzero, and where can't be used
- mask = attention_matrix < 1
- attention_matrix[mask] = attention.flatten()
- attention_matrix *= mask
- # 3 0 1 2
- # 0 3 1 2
- # 0 1 3 2
- # 0 1 2 3
+ # Use pre-computed nonzero results
+ non_diag_inds = self.pre_non_diag_inds.to(attention_matrix.device)
else:
non_diag_inds = torch.nonzero(attention_matrix == 0., as_tuple=False)
- attention_matrix[:] = 0
- attention_matrix[non_diag_inds[:, 0], non_diag_inds[:, 1], non_diag_inds[:, 2]] = attention.flatten()
+ attention_matrix[:] = 0
+ attention_matrix[non_diag_inds[:, 0], non_diag_inds[:, 1], non_diag_inds[:, 2]] = attention.flatten()
batch_anchor_features = batch_anchor_features.reshape(x.shape[0], len(self.anchors), -1)
attention_features = torch.bmm(torch.transpose(batch_anchor_features, 1, 2),
torch.transpose(attention_matrix, 1, 2)).transpose(1, 2)
diff --git a/utils/onnx_utils.py b/utils/onnx_utils.py
index a31b5933..7337380a 100644
--- a/utils/onnx_utils.py
+++ b/utils/onnx_utils.py
@@ -18,7 +18,8 @@
TRACE_REQUIRE_PREPROCESSING = [
'LSTR',
'RESA',
- 'RESA_Net'
+ 'RESA_Net',
+ 'LaneAtt'
]