make device configurable

Author: mmz33

i6_models/decoder/attention.py

@@ -39,18 +39,20 @@ def forward(
         query: torch.Tensor,
         weight_feedback: torch.Tensor,
         enc_seq_len: torch.Tensor,
+        device: str,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         :param key: encoder keys of shape [B,T,D_k]
         :param value: encoder values of shape [B,T,D_v]
         :param query: query of shape [B,D_k]
         :param weight_feedback: shape is [B,T,D_k]
         :param enc_seq_len: encoder sequence lengths [B]
+        :param device: device where to run the model (cpu or cuda)
         :return: attention context [B,D_v], attention weights [B,T,1]
         """
         # all inputs are already projected
         energies = self.linear(nn.functional.tanh(key + query.unsqueeze(1) + weight_feedback))  # [B,T,1]
-        time_arange = torch.arange(energies.size(1), device="cuda")  # [T]
+        time_arange = torch.arange(energies.size(1), device=device)  # [T]
         seq_len_mask = torch.less(time_arange[None, :], enc_seq_len[:, None])  # [B,T]
         energies = torch.where(seq_len_mask.unsqueeze(2), energies, torch.tensor(-float("inf")))
         weights = nn.functional.softmax(energies, dim=1)  # [B,T,1]
@@ -74,6 +76,7 @@ class AttentionLSTMDecoderV1Config:
         attention_cfg: attention config
         output_proj_dim: output projection dimension
         output_dropout: output dropout
+        device: device where to run the model (cpu or cuda)
     """
 
     encoder_dim: int
@@ -86,6 +89,7 @@ class AttentionLSTMDecoderV1Config:
     attention_cfg: AdditiveAttentionConfig
     output_proj_dim: int
     output_dropout: float
+    device: str
 
 
 class AttentionLSTMDecoderV1(nn.Module):
@@ -126,6 +130,8 @@ def __init__(self, cfg: AttentionLSTMDecoderV1Config):
         self.output = nn.Linear(cfg.output_proj_dim // 2, cfg.vocab_size)
         self.output_dropout = nn.Dropout(cfg.output_dropout)
 
+        self.device = cfg.device
+
     def forward(
         self,
         encoder_outputs: torch.Tensor,
@@ -140,10 +146,10 @@ def forward(
         :param state: decoder state
         """
         if state is None:
-            zeros = torch.zeros((encoder_outputs.size(0), self.lstm_hidden_size), device="cuda")
+            zeros = torch.zeros((encoder_outputs.size(0), self.lstm_hidden_size), device=self.device)
             lstm_state = (zeros, zeros)
-            att_context = torch.zeros((encoder_outputs.size(0), encoder_outputs.size(2)), device="cuda")
-            accum_att_weights = torch.zeros((encoder_outputs.size(0), encoder_outputs.size(1), 1), device="cuda")
+            att_context = torch.zeros((encoder_outputs.size(0), encoder_outputs.size(2)), device=self.device)
+            accum_att_weights = torch.zeros((encoder_outputs.size(0), encoder_outputs.size(1), 1), device=self.device)
         else:
             lstm_state, att_context, accum_att_weights = state
 
@@ -179,6 +185,7 @@ def forward(
                 query=s_transformed,
                 weight_feedback=weight_feedback,
                 enc_seq_len=enc_seq_len,
+                device=self.device,
             )
             att_context_list.append(att_context)
             accum_att_weights = accum_att_weights + att_weights * enc_inv_fertility * 0.5

tests/test_enc_dec_att.py

@@ -15,7 +15,9 @@ def test_additive_attention():
     enc_seq_len = torch.arange(start=10, end=20)  # [10, ..., 19]
 
     # pass key as weight feedback just for testing
-    context, weights = att(key=key, value=value, query=query, weight_feedback=key, enc_seq_len=enc_seq_len)
+    context, weights = att(
+        key=key, value=value, query=query, weight_feedback=key, enc_seq_len=enc_seq_len, device="cpu"
+    )
     assert context.shape == (10, 5)
     assert weights.shape == (10, 20, 1)
 
@@ -40,6 +42,7 @@ def test_encoder_decoder_attention_model():
         output_dropout=0.1,
         zoneout_drop_c=0.0,
         zoneout_drop_h=0.0,
+        device="cpu",
     )
     decoder = AttentionLSTMDecoderV1(decoder_cfg)
     target_labels = torch.randint(low=0, high=15, size=(10, 7))  # [B,N]
@@ -66,6 +69,7 @@ def forward_decoder(zoneout_drop_c: float, zoneout_drop_h: float):
             output_dropout=0.1,
             zoneout_drop_c=zoneout_drop_c,
             zoneout_drop_h=zoneout_drop_h,
+            device="cpu",
         )
         decoder = AttentionLSTMDecoderV1(decoder_cfg)
         decoder_logits, _ = decoder(encoder_outputs=encoder, labels=target_labels, enc_seq_len=encoder_seq_len)