""" PyTorch MCTS测试 测试统一的PyTorch MCTS实现 """ import pytest import torch import time import numpy as np from game import Game2048 from torch_mcts import TorchMCTS from training_data import TrainingDataManager class TestTorchMCTS: """PyTorch MCTS测试类""" @pytest.fixture def game(self): """测试游戏状态""" return Game2048(height=3, width=3, seed=42) @pytest.fixture def cpu_mcts(self): """CPU MCTS实例""" return TorchMCTS( c_param=1.414, max_simulation_depth=30, batch_size=1024, device="cpu" ) @pytest.fixture def gpu_mcts(self): """GPU MCTS实例""" if not torch.cuda.is_available(): pytest.skip("CUDA不可用") return TorchMCTS( c_param=1.414, max_simulation_depth=30, batch_size=4096, device="cuda" ) def test_cpu_mcts_basic_functionality(self, game, cpu_mcts): """测试CPU MCTS基本功能""" # 执行搜索 action, stats = cpu_mcts.search(game, 1000) # 验证结果 assert action in game.get_valid_moves(), f"选择了无效动作: {action}" assert 'action_visits' in stats assert 'action_avg_values' in stats assert 'sims_per_second' in stats assert stats['device'] == 'cpu' # 验证访问次数 total_visits = sum(stats['action_visits'].values()) assert total_visits == 1000, f"访问次数不匹配: {total_visits}" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA不可用") def test_gpu_mcts_basic_functionality(self, game, gpu_mcts): """测试GPU MCTS基本功能""" # 执行搜索 action, stats = gpu_mcts.search(game, 2000) # 验证结果 assert action in game.get_valid_moves(), f"选择了无效动作: {action}" assert 'action_visits' in stats assert 'action_avg_values' in stats assert 'sims_per_second' in stats assert stats['device'] == 'cuda' # 验证访问次数 total_visits = sum(stats['action_visits'].values()) assert total_visits == 2000, f"访问次数不匹配: {total_visits}" def test_action_distribution_quality(self, game, cpu_mcts): """测试动作分布质量""" action, stats = cpu_mcts.search(game, 5000) action_visits = stats['action_visits'] visit_values = list(action_visits.values()) # 检查分布不应该完全均匀（MCTS应该有偏向性） assert len(set(visit_values)) > 1, "动作分布完全均匀，不符合MCTS预期" # 检查最佳动作应该有最多访问次数 best_action_visits = action_visits[action] assert best_action_visits == max(visit_values), "最佳动作访问次数不是最多" # 检查价值的合理性 action_values = stats['action_avg_values'] for act, value in action_values.items(): assert value > 0, f"动作{act}的价值应该为正: {value}" assert value < 100000, f"动作{act}的价值过大: {value}" def test_device_auto_selection(self, game): """测试设备自动选择""" mcts = TorchMCTS(device="auto", batch_size=1024) # 验证设备选择 if torch.cuda.is_available(): assert mcts.device.type == "cuda" else: assert mcts.device.type == "cpu" # 执行搜索验证功能 action, stats = mcts.search(game, 1000) assert action in game.get_valid_moves() if mcts.device.type == "cuda": del mcts torch.cuda.empty_cache() def test_batch_size_auto_selection(self, game): """测试批次大小自动选择""" # CPU自动选择 cpu_mcts = TorchMCTS(device="cpu", batch_size=None) assert cpu_mcts.batch_size == 4096 # CPU默认批次大小 # GPU自动选择（如果可用） if torch.cuda.is_available(): gpu_mcts = TorchMCTS(device="cuda", batch_size=None) assert gpu_mcts.batch_size == 32768 # GPU默认批次大小 del gpu_mcts torch.cuda.empty_cache() def test_performance_cpu(self, game, cpu_mcts): """测试CPU性能""" simulations = 2000 start_time = time.time() action, stats = cpu_mcts.search(game, simulations) elapsed_time = time.time() - start_time speed = simulations / elapsed_time # CPU应该达到基本性能要求 assert speed > 100, f"CPU性能过低: {speed:.1f} 模拟/秒" # 验证统计信息准确性 assert abs(stats['sims_per_second'] - speed) < speed * 0.2, "统计信息不准确" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA不可用") def test_performance_gpu(self, game, gpu_mcts): """测试GPU性能""" simulations = 5000 torch.cuda.synchronize() start_time = time.time() action, stats = gpu_mcts.search(game, simulations) torch.cuda.synchronize() elapsed_time = time.time() - start_time speed = simulations / elapsed_time # GPU应该有合理的性能 assert speed > 50, f"GPU性能过低: {speed:.1f} 模拟/秒" # 验证统计信息准确性 assert abs(stats['sims_per_second'] - speed) < speed * 0.2, "统计信息不准确" def test_training_data_collection(self, game): """测试训练数据收集""" # 创建训练数据管理器 training_manager = TrainingDataManager( data_dir="data/test_torch_training", cache_size=5000, board_size=(3, 3) ) mcts = TorchMCTS( max_simulation_depth=30, batch_size=1024, device="cpu", training_manager=training_manager ) # 执行搜索 action, stats = mcts.search(game, 2000) # 验证训练数据收集 cache_stats = training_manager.get_cache_stats() assert cache_stats['cache_size'] > 0, "未收集到训练数据" # 验证数据质量 assert cache_stats['cache_size'] <= 2000, "收集的样本数超出预期" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA不可用") def test_memory_management(self, game): """测试GPU内存管理""" torch.cuda.empty_cache() initial_memory = torch.cuda.memory_allocated() gpu_mcts = TorchMCTS( max_simulation_depth=30, batch_size=8192, device="cuda" ) # 执行搜索 action, stats = gpu_mcts.search(game, 3000) # 检查内存使用 peak_memory = torch.cuda.max_memory_allocated() memory_used = (peak_memory - initial_memory) / 1e6 # MB assert memory_used < 200, f"GPU内存使用过多: {memory_used:.1f} MB" # 清理并验证内存释放 del gpu_mcts torch.cuda.empty_cache() final_memory = torch.cuda.memory_allocated() assert final_memory <= initial_memory * 1.1, "GPU内存未正确释放" def test_device_switching(self, game): """测试设备动态切换""" mcts = TorchMCTS(device="cpu", batch_size=1024) # 初始为CPU assert mcts.device.type == "cpu" action1, stats1 = mcts.search(game.copy(), 1000) assert stats1['device'] == 'cpu' # 切换到GPU（如果可用） if torch.cuda.is_available(): mcts.set_device("cuda") assert mcts.device.type == "cuda" action2, stats2 = mcts.search(game.copy(), 1000) assert stats2['device'] == 'cuda' # 切换回CPU mcts.set_device("cpu") assert mcts.device.type == "cpu" torch.cuda.empty_cache() def test_consistency_across_devices(self, game): """测试不同设备间的一致性""" if not torch.cuda.is_available(): pytest.skip("CUDA不可用") # 使用相同的随机种子 np.random.seed(42) cpu_mcts = TorchMCTS(device="cpu", batch_size=2048) cpu_action, cpu_stats = cpu_mcts.search(game.copy(), 3000) np.random.seed(42) gpu_mcts = TorchMCTS(device="cuda", batch_size=2048) gpu_action, gpu_stats = gpu_mcts.search(game.copy(), 3000) # 由于随机性，动作可能不完全一致，但应该在合理范围内 # 这里主要验证两个设备都能正常工作 assert cpu_action in game.get_valid_moves() assert gpu_action in game.get_valid_moves() # 验证访问次数总和 cpu_total = sum(cpu_stats['action_visits'].values()) gpu_total = sum(gpu_stats['action_visits'].values()) assert cpu_total == gpu_total == 3000 del gpu_mcts torch.cuda.empty_cache() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA不可用") def test_batch_size_optimization(): """测试批次大小优化""" game = Game2048(height=3, width=3, seed=42) mcts = TorchMCTS(device="cuda", batch_size=4096) # 执行批次大小优化 optimal_size = mcts.optimize_batch_size(game, test_simulations=1000) # 验证优化结果 assert optimal_size > 0 assert mcts.batch_size == optimal_size # 验证优化后的性能 action, stats = mcts.search(game, 2000) assert action in game.get_valid_moves() assert stats['sims_per_second'] > 0 del mcts torch.cuda.empty_cache() if __name__ == "__main__": pytest.main([__file__, "-v"])