Fix global PyTorch initializer leak in build_vae_var

models/__init__.py

@@ -20,20 +20,27 @@ def build_vae_var(
     width = depth * 64
     dpr = 0.1 * depth/24
 
-    # disable built-in initialization for speed
-    for clz in (nn.Linear, nn.LayerNorm, nn.BatchNorm2d, nn.SyncBatchNorm, nn.Conv1d, nn.Conv2d, nn.ConvTranspose1d, nn.ConvTranspose2d):
-        setattr(clz, 'reset_parameters', lambda self: None)
-
-    # build models
-    vae_local = VQVAE(vocab_size=V, z_channels=Cvae, ch=ch, test_mode=True, share_quant_resi=share_quant_resi, v_patch_nums=patch_nums).to(device)
-    var_wo_ddp = VAR(
-        vae_local=vae_local,
-        num_classes=num_classes, depth=depth, embed_dim=width, num_heads=heads, drop_rate=0., attn_drop_rate=0., drop_path_rate=dpr,
-        norm_eps=1e-6, shared_aln=shared_aln, cond_drop_rate=0.1,
-        attn_l2_norm=attn_l2_norm,
-        patch_nums=patch_nums,
-        flash_if_available=flash_if_available, fused_if_available=fused_if_available,
-    ).to(device)
+    # disable built-in initialization for speed while constructing these models
+    init_classes = (nn.Linear, nn.LayerNorm, nn.BatchNorm2d, nn.SyncBatchNorm, nn.Conv1d, nn.Conv2d, nn.ConvTranspose1d, nn.ConvTranspose2d)
+    reset_parameters = {clz: clz.reset_parameters for clz in init_classes}
+    try:
+        for clz in init_classes:
+            setattr(clz, 'reset_parameters', lambda self: None)
+
+        # build models
+        vae_local = VQVAE(vocab_size=V, z_channels=Cvae, ch=ch, test_mode=True, share_quant_resi=share_quant_resi, v_patch_nums=patch_nums).to(device)
+        var_wo_ddp = VAR(
+            vae_local=vae_local,
+            num_classes=num_classes, depth=depth, embed_dim=width, num_heads=heads, drop_rate=0., attn_drop_rate=0., drop_path_rate=dpr,
+            norm_eps=1e-6, shared_aln=shared_aln, cond_drop_rate=0.1,
+            attn_l2_norm=attn_l2_norm,
+            patch_nums=patch_nums,
+            flash_if_available=flash_if_available, fused_if_available=fused_if_available,
+        ).to(device)
+    finally:
+        for clz, reset_parameter in reset_parameters.items():
+            setattr(clz, 'reset_parameters', reset_parameter)
+
     var_wo_ddp.init_weights(init_adaln=init_adaln, init_adaln_gamma=init_adaln_gamma, init_head=init_head, init_std=init_std)
 
     return vae_local, var_wo_ddp

tests/test_model_initialization.py

@@ -0,0 +1,60 @@
+import unittest
+from unittest.mock import patch
+
+import torch.nn as nn
+
+import models
+
+
+INIT_CLASSES = (
+    nn.Linear,
+    nn.LayerNorm,
+    nn.BatchNorm2d,
+    nn.SyncBatchNorm,
+    nn.Conv1d,
+    nn.Conv2d,
+    nn.ConvTranspose1d,
+    nn.ConvTranspose2d,
+)
+
+
+class ModelInitializationTests(unittest.TestCase):
+    def setUp(self):
+        self.reset_parameters = {clz: clz.reset_parameters for clz in INIT_CLASSES}
+
+    def tearDown(self):
+        for clz, reset_parameter in self.reset_parameters.items():
+            clz.reset_parameters = reset_parameter
+
+    def assert_initializers_restored(self):
+        for clz, reset_parameter in self.reset_parameters.items():
+            with self.subTest(module=clz.__name__):
+                self.assertIs(clz.reset_parameters, reset_parameter)
+
+    def test_build_vae_var_restores_builtin_initializers(self):
+        models.build_vae_var(
+            device="cpu",
+            patch_nums=(1, 2),
+            V=16,
+            Cvae=4,
+            ch=32,
+            num_classes=2,
+            depth=1,
+            flash_if_available=False,
+            fused_if_available=False,
+        )
+
+        self.assert_initializers_restored()
+
+    def test_build_vae_var_restores_initializers_after_failure(self):
+        with patch.object(
+            models, "VQVAE", side_effect=RuntimeError("construction failed")
+        ):
+            with self.assertRaisesRegex(RuntimeError, "construction failed"):
+                models.build_vae_var(device="cpu")
+
+        self.assert_initializers_restored()
+
+
+if __name__ == "__main__":
+    unittest.main()