Optimize exposure autoencoder distributed training

2026-07-09 15:44:21 +08:00
parent f7fb6b7718
commit 5a3122c965
3 changed files with 258 additions and 50 deletions
--- a/backbones.py
+++ b/backbones.py
@@ -232,7 +232,8 @@ class TimesNetBlock(nn.Module):
        amplitude[0] = 0.0
        k = min(self.top_k, amplitude.numel() - 1)
        weights, indices = torch.topk(amplitude, k=k)
-        periods = [max(1, T // int(idx.item())) for idx in indices]
+        # One device synchronization per block instead of one per selected period.
+        periods = [max(1, T // int(idx)) for idx in indices.tolist()]
        return periods, weights.to(dtype=x.dtype, device=x.device)

    def _period_branch(self, x: torch.Tensor, period: int) -> torch.Tensor: