Skip to content

Commit fdc839a

Browse files
authored
perf: avoid graph break for SiLUT when inferring (#4790)
This pull request simplifies and optimizes the implementation of the `forward` method in the `ActivationFn` class within `deepmd/pt/utils/utils.py`. The changes streamline the logic by removing unnecessary condition checks and directly using `torch.where` for computation. I've evaluated this change using inference efficiency tasks from LAMBench with DPA 3.1 3M model. | System | Before: Avg Time ± Std (s) | After: Avg Time ± Std (s) | Speedup | Success Rate | |-------------------|---------------------------|--------------------------|---------|--------------| | `catalysts_500.traj` | 211.82 ± 19.31 | **196.14 ± 18.11** | +7.1% | 100.0% | | `inorganic_500.traj` | 204.62 ± 40.22 | **191.20 ± 36.44** | +6.4% | 100.0% | <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Refactor** - Improved the internal logic of the SiLU activation function for more streamlined processing. No changes to user-facing functionality. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent c46dc7d commit fdc839a

File tree

1 file changed

+4
-7
lines changed

1 file changed

+4
-7
lines changed

deepmd/pt/utils/utils.py

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -149,13 +149,10 @@ def silu_grad(x):
149149
self.const = float(silu(threshold))
150150

151151
def forward(self, x: torch.Tensor) -> torch.Tensor:
152-
silu_part = F.silu(x)
153-
mask = x >= self.threshold
154-
if torch.any(mask):
155-
tanh_part = torch.tanh(self.slope * (x - self.threshold)) + self.const
156-
return torch.where(x < self.threshold, silu_part, tanh_part)
157-
else:
158-
return silu_part
152+
sig = torch.sigmoid(x)
153+
silu = x * sig
154+
tanh = torch.tanh(self.slope * (x - self.threshold)) + self.const
155+
return torch.where(x >= self.threshold, tanh, silu)
159156

160157

161158
class ActivationFn(torch.nn.Module):

0 commit comments

Comments
 (0)