Update pretrain_inference.py

pretrain_inference.py

@@ -1,131 +1,160 @@
-import torch
-import torch.nn.functional as F
-from tokenizers import Tokenizer
-import intel_extension_for_pytorch as ipex
-from novel_model import NovelTransformer
-
-# 配置参数
-VOCAB_SIZE = 8000
-D_MODEL = 128
-NHEAD = 4
-NUM_LAYERS = 4
-DIM_FEEDFORWARD = 512
-DROPOUT = 0.1
-MAX_LEN = 4096
-MODEL_PATH = "d:/图像/novel_model/best_model.pt"
-TOKENIZER_PATH = "d:/图像/novel_tokenizer.json"
-
-def generate_text(model, tokenizer, prompt, max_length=100, temperature=0.8, top_k=50, top_p=0.9, device="cpu"):
-    """生成文本"""
-    model.eval()
-    
-    # 编码提示
-    input_ids = torch.tensor(tokenizer.encode(prompt).ids, dtype=torch.long).unsqueeze(0).to(device)
-    
-    # 生成文本
-    with torch.no_grad():
-        for _ in range(max_length):
-            # 如果序列太长，截断
-            if input_ids.size(1) > MAX_LEN:
-                input_ids = input_ids[:, -MAX_LEN:]
-            
-            # 获取模型输出
-            outputs = model(input_ids)
-            next_token_logits = outputs[:, -1, :] / temperature
-            
-            # 应用top-k过滤
-            if top_k > 0:
-                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
-                next_token_logits[indices_to_remove] = float('-inf')
-            
-            # 应用top-p过滤
-            if top_p < 1.0:
-                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
-                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-                
-                # 移除概率累积超过阈值的token
-                sorted_indices_to_remove = cumulative_probs > top_p
-                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-                sorted_indices_to_remove[..., 0] = 0
-                
-                indices_to_remove = sorted_indices[sorted_indices_to_remove]
-                next_token_logits[0, indices_to_remove] = float('-inf')
-            
-            # 采样下一个token
-            probs = F.softmax(next_token_logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-            
-            # 添加到输入序列
-            input_ids = torch.cat([input_ids, next_token], dim=1)
-            
-            # 如果生成了结束标记，停止生成
-            if next_token.item() == tokenizer.token_to_id("</s>"):
-                break
-    
-    # 解码生成的ID
-    output = tokenizer.decode(input_ids[0].tolist())
-    return output
-
-def main():
-    # 设置设备
-    device = torch.device("xpu" if torch.xpu.is_available() else "cpu")
-    print(f"使用设备: {device}")
-    
-    # 加载分词器
-    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
-    
-    # 加载模型
-    checkpoint = torch.load(MODEL_PATH, map_location=device)
-    
-    model = NovelTransformer(
-        vocab_size=VOCAB_SIZE,
-        d_model=D_MODEL,
-        nhead=NHEAD,
-        num_layers=NUM_LAYERS,
-        dim_feedforward=DIM_FEEDFORWARD,
-        dropout=DROPOUT,
-        max_len=MAX_LEN
-    )
-    
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model = model.to(device)
-    
-    # 设置为评估模式
-    model.eval()
-    
-    # 对于推理，不传递优化器参数
-    # 在较旧版本的IPEX中，可以使用以下方式优化推理
-    with torch.no_grad():
-        model = ipex.optimize(model)
-    
-    # 交互式生成
-    print("预训练小说语言模型已加载。输入提示进行生成，输入'exit'退出。")
-    
-    # 测试模型困惑度
-    test_text = "从前有座山，山上有座庙，庙里有个"
-    print(f"\n测试文本: {test_text}")
-    
-    # 生成文本
-    generated = generate_text(model, tokenizer, test_text, max_length=50, device=device)
-    print("\n生成的文本:")
-    print(generated)
-    
-    # 交互式生成
-    while True:
-        prompt = input("\n请输入提示 (或输入'exit'退出): ")
-        if prompt.lower() == 'exit':
-            break
-        
-        # 生成参数
-        length = int(input("生成长度 (默认100): ") or "100")
-        temp = float(input("温度 (0.1-1.0, 默认0.8): ") or "0.8")
-        
-        # 生成文本
-        generated = generate_text(model, tokenizer, prompt, max_length=length, 
-                                 temperature=temp, device=device)
-        
-        print("\n生成的文本:")
-        print(generated)
-
-if __name__ == "__main__":
+import torch
+import torch.nn.functional as F
+from tokenizers import Tokenizer
+from novel_model import NovelTransformer
+
+# 配置参数
+VOCAB_SIZE = 8000
+D_MODEL = 128
+NHEAD = 4
+NUM_LAYERS = 4
+DIM_FEEDFORWARD = 512
+DROPOUT = 0.1
+MAX_LEN = 4096
+MODEL_PATH = "d:/图像/novel_model/best_model.pt"
+TOKENIZER_PATH = "d:/图像/novel_tokenizer.json"
+
+def generate_text(model, tokenizer, prompt, max_length=100, temperature=0.8, top_k=50, top_p=0.9, device="cuda"):
+    """生成文本"""
+    model.eval()
+    
+    # 编码提示
+    input_ids = torch.tensor(tokenizer.encode(prompt).ids, dtype=torch.long).unsqueeze(0).to(device)
+    
+    # 生成文本
+    with torch.no_grad():
+        for _ in range(max_length):
+            # 如果序列太长，截断
+            if input_ids.size(1) > MAX_LEN:
+                input_ids = input_ids[:, -MAX_LEN:]
+            
+            # 获取模型输出
+            outputs = model(input_ids)
+            next_token_logits = outputs[:, -1, :] / temperature
+            
+            # 应用top-k过滤
+            if top_k > 0:
+                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                next_token_logits[indices_to_remove] = float('-inf')
+            
+            # 应用top-p过滤
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                
+                # 移除概率累积超过阈值的token
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+                
+                indices_to_remove = sorted_indices[sorted_indices_to_remove]
+                next_token_logits[0, indices_to_remove] = float('-inf')
+            
+            # 采样下一个token
+            probs = F.softmax(next_token_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            
+            # 添加到输入序列
+            input_ids = torch.cat([input_ids, next_token], dim=1)
+            
+            # 如果生成了结束标记，停止生成
+            if next_token.item() == tokenizer.token_to_id("</s>"):
+                break
+    
+    # 解码生成的ID
+    output = tokenizer.decode(input_ids[0].tolist())
+    return output
+
+def main():
+    # 设置设备
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"使用设备: {device}")
+    
+    # 显示CUDA信息
+    if device.type == 'cuda':
+        print(f"CUDA设备: {torch.cuda.get_device_name(0)}")
+        print(f"CUDA版本: {torch.version.cuda}")
+        print(f"当前GPU内存使用: {torch.cuda.memory_allocated(0)/1024**2:.2f} MB")
+    
+    # 加载分词器
+    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
+    
+    # 加载模型
+    checkpoint = torch.load(MODEL_PATH, map_location=device)
+    
+    model = NovelTransformer(
+        vocab_size=VOCAB_SIZE,
+        d_model=D_MODEL,
+        nhead=NHEAD,
+        num_layers=NUM_LAYERS,
+        dim_feedforward=DIM_FEEDFORWARD,
+        dropout=DROPOUT,
+        max_len=MAX_LEN
+    )
+    
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model = model.to(device)
+    
+    # 设置为评估模式
+    model.eval()
+    
+    # 启用CUDA优化
+    if device.type == 'cuda':
+        torch.backends.cudnn.benchmark = True
+    
+    # 测试模型困惑度
+    test_text = "从前有座山，山上有座庙，庙里有个"
+    print(f"\n测试文本: {test_text}")
+    
+    # 生成文本
+    generated = generate_text(model, tokenizer, test_text, max_length=50, device=device)
+    print("\n生成的文本:")
+    print(generated)
+    
+    # 交互式生成
+    print("\n预训练小说语言模型已加载。输入提示进行生成，输入'exit'退出。")
+    
+    while True:
+        prompt = input("\n请输入提示 (或输入'exit'退出): ")
+        if prompt.lower() == 'exit':
+            break
+        
+        # 生成参数
+        length = int(input("生成长度 (默认100): ") or "100")
+        temp = float(input("温度 (0.1-1.0, 默认0.8): ") or "0.8")
+        top_k_val = int(input("Top-K (默认50): ") or "50")
+        top_p_val = float(input("Top-P (0.0-1.0, 默认0.9): ") or "0.9")
+        
+        # 记录生成开始时间
+        start_time = torch.cuda.Event(enable_timing=True)
+        end_time = torch.cuda.Event(enable_timing=True)
+        
+        start_time.record()
+        
+        # 生成文本
+        generated = generate_text(
+            model, 
+            tokenizer, 
+            prompt, 
+            max_length=length, 
+            temperature=temp,
+            top_k=top_k_val,
+            top_p=top_p_val,
+            device=device
+        )
+        
+        end_time.record()
+        
+        # 等待CUDA操作完成
+        torch.cuda.synchronize()
+        
+        # 计算生成时间
+        generation_time = start_time.elapsed_time(end_time) / 1000  # 转换为秒
+        
+        print("\n生成的文本:")
+        print(generated)
+        print(f"\n生成时间: {generation_time:.2f} 秒")
+        print(f"生成速度: {length/generation_time:.2f} 字符/秒")
+
+if __name__ == "__main__":
     main()