Merge pull request lucidrains#206 from LWprogramming/inference_mode

Change torch.no_grad() to torch.inference_mode()

Author: lucidrains

audiolm_pytorch/audiolm_pytorch.py

@@ -531,7 +531,7 @@ def forward(
 
         text_mask = None
         if not exists(text_embeds) and exists(text):
-            with torch.no_grad():
+            with torch.inference_mode():
                 text_embeds = self.embed_text(text, output_device = device)
                 text_mask = torch.any(text_embeds != 0, dim = -1)
 
@@ -677,7 +677,7 @@ def forward(
         assert not (self.has_condition ^ has_text)
 
         if not exists(text_embeds) and exists(text):
-            with torch.no_grad():
+            with torch.inference_mode():
                 text_embeds = self.embed_text(text, output_device = device)
 
         text_mask = None
@@ -907,7 +907,7 @@ def forward(
 
         text_mask = None
         if not exists(text_embeds) and exists(text):
-            with torch.no_grad():
+            with torch.inference_mode():
                 text_embeds = self.embed_text(text, output_device = device)
                 text_mask = torch.any(text_embeds != 0, dim = -1)
 
@@ -1141,7 +1141,7 @@ def embed_text(self, text):
         return self.transformer.embed_text(text, output_device = self.device)
 
     @eval_decorator
-    @torch.no_grad()
+    @torch.inference_mode()
     @beartype
     def generate(
         self,
@@ -1186,7 +1186,7 @@ def generate(
         assert not (self.transformer.has_condition ^ has_text)
 
         if not exists(text_embeds) and exists(text):
-            with torch.no_grad():
+            with torch.inference_mode():
                 text_embeds = self.transformer.embed_text(text, output_device = device)
 
         # start length and get running id output
@@ -1323,7 +1323,7 @@ def device(self):
         return next(self.parameters()).device
 
     @eval_decorator
-    @torch.no_grad()
+    @torch.inference_mode()
     @beartype
     def generate(
         self,
@@ -1353,7 +1353,7 @@ def generate(
             coarse_token_ids = prime_coarse_token_ids
         elif exists(prime_wave):
             assert exists(self.codec)
-            with torch.no_grad():
+            with torch.inference_mode():
                 self.codec.eval()
                 _, indices, _ = self.codec(prime_wave, return_encoded = True)
                 coarse_token_ids = indices[..., :self.num_coarse_quantizers]
@@ -1367,7 +1367,7 @@ def generate(
         assert not (self.transformer.has_condition ^ has_text)
 
         if not exists(text_embeds) and exists(text):
-            with torch.no_grad():
+            with torch.inference_mode():
                 text_embeds = self.transformer.embed_text(text, output_device = device)
 
         if self.unique_consecutive:
@@ -1444,7 +1444,7 @@ def forward(
         if not exists(coarse_token_ids):
             assert exists(self.codec), 'Codec must be provided if given raw wave for training'
 
-            with torch.no_grad():
+            with torch.inference_mode():
                 self.codec.eval()
                 _, indices, _ = self.codec(raw_wave_for_codec, return_encoded = True)
 
@@ -1568,7 +1568,7 @@ def device(self):
         return next(self.parameters()).device
 
     @eval_decorator
-    @torch.no_grad()
+    @torch.inference_mode()
     @beartype
     def generate(
         self,
@@ -1597,7 +1597,7 @@ def generate(
         assert not (self.transformer.has_condition ^ has_text)
 
         if not exists(text_embeds) and exists(text):
-            with torch.no_grad():
+            with torch.inference_mode():
                 text_embeds = self.transformer.embed_text(text, output_device = device)
 
         # initialize fine token ids
@@ -1609,7 +1609,7 @@ def generate(
             fine_token_ids = prime_fine_token_ids
         elif exists(prime_wave):
             assert exists(self.codec)
-            with torch.no_grad():
+            with torch.inference_mode():
                 self.codec.eval()
                 _, token_ids, _ = self.codec(prime_wave, return_encoded = True)
 
@@ -1698,7 +1698,7 @@ def forward(
         if exists(raw_wave):
             assert exists(self.codec), 'Codec must be provided if given raw wave for training'
 
-            with torch.no_grad():
+            with torch.inference_mode():
                 self.codec.eval()
                 _, token_ids, _ = self.codec(raw_wave, return_encoded = True)
 
@@ -1829,7 +1829,7 @@ def device(self):
         return next(self.parameters()).device
 
     @eval_decorator
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         *,

audiolm_pytorch/encodec.py

@@ -89,7 +89,7 @@ def forward(
         wav = rearrange(x, f'b t -> b {self.model.channels} t')
 
         # Extract discrete codes from EnCodec
-        with torch.no_grad():
+        with torch.inference_mode():
             encoded_frames = self.model.encode(wav)
         # encoded_frames is a list of (frame, scale) tuples. Scale is a scalar but we don't use it. Frame is a tensor
         # of shape [batch, num_quantizers, num_samples_per_frame]. We want to concatenate the frames to get all the

audiolm_pytorch/hubert_kmeans.py

@@ -75,7 +75,7 @@ def downsample_factor(self):
         # todo: double check
         return 320
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         wav_input,

audiolm_pytorch/t5.py

@@ -93,7 +93,7 @@ def t5_encode_text(
 
     t5.eval()
 
-    with torch.no_grad():
+    with torch.inference_mode():
         output = t5(input_ids = input_ids, attention_mask = attn_mask)
         encoded_text = output.last_hidden_state.detach()
 

audiolm_pytorch/trainer.py

@@ -507,7 +507,7 @@ def train_step(self):
             for model, label in models:
                 model.eval()
 
-                with torch.no_grad():
+                with torch.inference_mode():
                     recons = model(wave, return_recons_only = True)
 
                 for ind, recon in enumerate(recons.unbind(dim = 0)):
@@ -753,7 +753,7 @@ def train_step(self):
         if self.is_main and not (steps % self.save_results_every):
             data_kwargs = self.data_tuple_to_kwargs(next(self.valid_dl_iter))
 
-            with torch.no_grad():
+            with torch.inference_mode():
                 self.train_wrapper.eval()
                 valid_loss = self.train_wrapper(**data_kwargs, return_loss = True)
 
@@ -1002,7 +1002,7 @@ def train_step(self):
         if self.is_main and not (steps % self.save_results_every):
             data_kwargs = dict(zip(self.ds_fields, next(self.valid_dl_iter)))
 
-            with torch.no_grad():
+            with torch.inference_mode():
                 self.train_wrapper.eval()
 
                 valid_loss = self.train_wrapper(
@@ -1252,7 +1252,7 @@ def train_step(self):
         if self.is_main and not (steps % self.save_results_every):
             data_kwargs = self.data_tuple_to_kwargs(next(self.valid_dl_iter))
 
-            with torch.no_grad():
+            with torch.inference_mode():
                 self.train_wrapper.eval()
                 valid_loss = self.train_wrapper(**data_kwargs, return_loss = True)
 

audiolm_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.2.11'
+__version__ = '1.2.12'

audiolm_pytorch/vq_wav2vec.py

@@ -59,7 +59,7 @@ def downsample_factor(self):
     def codebook_size(self):
         return self.model.vector_quantizer.embedding.shape[0]
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         wav_input,