Changes to be committed: 한번에 생성하는 코드 중간에 남기기

Models/Vector2MIDI.py

@@ -1,6 +1,7 @@
-from torch import tanh, zeros, no_grad, full_like, topk, multinomial, cat, int64, nn, stack
 from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
 import torch.nn.functional as F
+from torch import nn
+import torch
 
 class Vector2MIDI(nn.Module):
     def __init__(self, hidden_dim, input_dim=25, dropout=0.2):
@@ -31,8 +32,8 @@ class Vector2MIDI(nn.Module):
 
     def init_hidden_states(self, x):
         """초기 hidden과 cell state 생성"""
-        h0 = tanh(self.init_hidden(x)) # 활성화 함수 추가 (hyperbolic tangent)
-        c0 = tanh(self.init_cell(x))
+        h0 = torch.tanh(self.init_hidden(x)) # 활성화 함수 추가 (hyperbolic tangent)
+        c0 = torch.tanh(self.init_cell(x))
         
         h0 = h0.unsqueeze(0).repeat(2, 1, 1)  # (num_layers, B, H)
         c0 = c0.unsqueeze(0).repeat(2, 1, 1)
@@ -52,11 +53,11 @@ class Vector2MIDI(nn.Module):
             embeddings.append(dim_onehot)
             
         # 모든 차원을 연결
-        full_embedding = cat(embeddings, dim=-1)  # (B, T, total_vocab)
+        full_embedding = torch.cat(embeddings, dim=-1)  # (B, T, total_vocab)
         
         # 스타일 컨텍스트를 각 타임스텝에 추가
         style_expanded = style_context.unsqueeze(1).expand(-1, seq_len, -1)  # (B, T, hidden_dim//2)
-        combined_input = cat([full_embedding, style_expanded], dim=-1)  # (B, T, total_vocab + hidden_dim//2)
+        combined_input = torch.cat([full_embedding, style_expanded], dim=-1)  # (B, T, total_vocab + hidden_dim//2)
         
         return self.input_embedding(combined_input)
 
@@ -85,68 +86,37 @@ class Vector2MIDI(nn.Module):
             outputs.append(dim_logits)
 
         return outputs
+    
+    def top_k_filtering(self, logits, top_k):
+        """Top-k 필터링"""
+        if top_k > 0:
+            # logits의 마지막 차원에서만 top-k 선택
+            values, _ = torch.topk(logits, min(top_k, logits.size(-1)), dim=-1)
+            min_values = values[..., -1:]  # 마지막 k번째 값, 모든 차원 유지
+            logits = torch.where(logits < min_values, torch.full_like(logits, float('-inf')), logits)
+        return logits
 
-    def generate(self, x, device, max_steps=1024, temperature:float=1.0, top_k=None):
-        self.eval()
+    def generate(self, x, device, seq_len=64, temperature:float=1.2, top_k=5):
+        self.eval() # autogressive로 한 타임 한 타임 생성하는 거 말고, forward를 이용해서 한 번에 생성하기로 변경
         x = x.to(device)
         batch_size = x.size(0)
-        
-        h, c = self.init_hidden_states(x)
-        style_context = self.style_context(x)  # (B, hidden_dim//2)
-        current_tokens = self.generate_start_tokens_from_style(x) # 첫 토큰
-        generated_tokens = zeros(batch_size, max_steps, 7, dtype=int64, device=device) # 생성될 토큰 저장 Tenosr
-        
-        with no_grad():
-            for step in range(max_steps):
-                # 현재 토큰을 임베딩으로 변환
-                embedded = self.tokens_to_embedding(current_tokens, style_context)
-                lstm_out, (h, c) = self.lstm(embedded, (h, c))  # lstm_out: (B,1,H)
-                hidden = self.fc_mid(lstm_out[:, -1, :])  # (B, 256)
-                
-                # 각 차원별로 다음 토큰 생성
-                next_tokens = []
-                for head in self.output_heads:
-                    logits = head(hidden)  # (B, vocab_size_i)
-                    
-                    if temperature != 1.0:
-                        logits = logits / temperature
-                    
-                    if top_k is not None and top_k > 0:
-                        k = min(top_k, logits.size(-1))
-                        topk_vals, topk_idx = topk(logits, k, dim=-1)
-                        # create mask with very low values
-                        low_val = -1e9
-                        mask = full_like(logits, low_val)
-                        logits = mask.scatter(-1, topk_idx, topk_vals)
 
-                    probs = F.softmax(logits, dim=-1)
-                    token = multinomial(probs, num_samples=1)  # (B,1)
-                    next_tokens.append(token)
-                
-                current_tokens = cat(next_tokens, dim=1).unsqueeze(1)  # (B, 1, 7)
-                generated_tokens[:, step, :] = current_tokens.squeeze(1)
-                
-                # 종료 조건
-                if (current_tokens == -1).all():
-                    break
-        
-        return generated_tokens
-    
-    def generate_start_tokens_from_style(self, x):
-        """스타일 벡터에서 첫 토큰 생성"""
-        batch_size = x.size(0)
-        start_tokens = zeros(batch_size, 1, 7, dtype=int64, device=x.device)
-        
-        for i, head in enumerate(self.start_token_heads):
-            logits = head(x)  # (B, vocab_size_i)
-            
-            # 스타일 기반 첫 토큰 샘플링
-            if i in [1, 4, 6]:  # duration 차원: 더 확정적으로
-                probs = F.softmax(logits / 0.5, dim=-1)  # 낮은 온도
-            else:  # pitch, velocity 등: 다양성 허용
-                probs = F.softmax(logits / 1.2, dim=-1)  # 약간 높은 온도
-            
-            token = multinomial(probs, num_samples=1)  # (B, 1)
-            start_tokens[:, :, i] = token
+        generated_sequence = torch.zeros((batch_size, seq_len, 7), dtype=torch.long, device=device)
+        lengths = torch.full((batch_size,), seq_len, dtype=torch.long, device=device)
+
+        with torch.no_grad():
+            logits_list = self.forward(x, lengths, generated_sequence)
+
+        # 첫 번째 토큰 이후부터 샘플링
+        for i, logits in enumerate(logits_list):
+            dim_logits = logits[:, :-1, :] / temperature  # (B, T-1, vocab)
+            dim_logits = self.top_k_filtering(dim_logits, top_k)
+            probs = F.softmax(dim_logits, dim=-1)
             
-        return start_tokens
+            # 각 타임스텝별로 샘플링
+            for t in range(seq_len):
+                if t < dim_logits.size(1):
+                    sampled = torch.multinomial(probs[:, t, :], 1).squeeze(-1)
+                    generated_sequence[:, t, i] = sampled
+
+        return generated_sequence