Add new file

jax/test.py

+import jax
+import jax.numpy as jnp
+from jax import random, grad, jit, vmap
+import time
+
+# 检查JAX是否使用了GPU
+print("JAX版本:", jax.__version__)
+print("可用设备:", jax.devices())
+print("当前使用的设备:", jax.devices()[0])
+print("设备类型:", jax.devices()[0].device_kind)
+
+# 1. 基本的GPU矩阵运算示例
+def basic_gpu_operations():
+    print("\n=== 基础GPU矩阵运算 ===")
+    
+    # 创建大型矩阵（这些会自动在GPU上分配）
+    key = random.PRNGKey(42)
+    size = 5000  # 大型矩阵
+    
+    # 生成随机矩阵
+    A = random.normal(key, (size, size))
+    B = random.normal(key, (size, size))
+    
+    # GPU矩阵乘法
+    start = time.time()
+    C = jnp.dot(A, B)  # 这会自动在GPU上执行
+    compute_time = time.time() - start
+    
+    print(f"矩阵大小: {size}x{size}")
+    print(f"GPU矩阵乘法耗时: {compute_time:.4f}秒")
+    
+    # 检查结果
+    print(f"结果矩阵C的形状: {C.shape}")
+    print(f"C的平均值: {jnp.mean(C):.6f}")
+    print(f"C的标准差: {jnp.std(C):.6f}")
+    
+    return C
+
+# 2. 使用JIT编译加速
+def jit_acceleration():
+    print("\n=== JIT编译加速示例 ===")
+    
+    # 定义一个复杂函数
+    def complex_function(x):
+        return jnp.sum(jnp.sin(x) * jnp.exp(-x**2) + jnp.log1p(jnp.abs(x)))
+    
+    # 创建输入数据
+    key = random.PRNGKey(0)
+    x = random.normal(key, (10000, 1000))
+    
+    # 不使用JIT
+    start = time.time()
+    result_no_jit = complex_function(x)
+    time_no_jit = time.time() - start
+    
+    # 使用JIT编译
+    jitted_function = jit(complex_function)
+    
+    # 第一次调用会有编译开销
+    start = time.time()
+    result_jit = jitted_function(x)
+    time_first_call = time.time() - start
+    
+    # 第二次调用（已经编译）
+    start = time.time()
+    result_jit = jitted_function(x)
+    time_second_call = time.time() - start
+    
+    print(f"不使用JIT: {time_no_jit:.4f}秒")
+    print(f"JIT首次调用（含编译）: {time_first_call:.4f}秒")
+    print(f"JIT后续调用: {time_second_call:.4f}秒")
+    print(f"加速比: {time_no_jit/time_second_call:.2f}倍")
+    
+    return result_jit
+
+# 3. 自动微分和GPU计算
+def autodiff_gpu():
+    print("\n=== GPU上的自动微分 ===")
+    
+    # 定义损失函数
+    def loss(params, x, y):
+        w, b = params
+        predictions = jnp.dot(x, w) + b
+        return jnp.mean((predictions - y)**2)
+    
+    # 生成模拟数据
+    key = random.PRNGKey(1)
+    key1, key2, key3 = random.split(key, 3)
+    
+    n_samples = 10000
+    n_features = 100
+    
+    X = random.normal(key1, (n_samples, n_features))
+    true_w = random.normal(key2, (n_features, 1))
+    true_b = random.normal(key3, (1,))
+    Y = jnp.dot(X, true_w) + true_b + 0.1 * random.normal(key1, (n_samples, 1))
+    
+    # 初始化参数
+    params = (random.normal(key1, (n_features, 1)), jnp.zeros((1,)))
+    
+    # 计算梯度的GPU函数
+    grad_loss = jit(grad(loss, argnums=0))
+    
+    # 在GPU上计算梯度
+    start = time.time()
+    gradients = grad_loss(params, X, Y)
+    grad_time = time.time() - start
+    
+    print(f"数据形状: X={X.shape}, Y={Y.shape}")
+    print(f"GPU梯度计算耗时: {grad_time:.4f}秒")
+    print(f"权重梯度形状: {gradients[0].shape}")
+    print(f"偏置梯度形状: {gradients[1].shape}")
+    
+    return gradients
+
+# 4. 批量处理示例
+def batch_processing():
+    print("\n=== GPU批量处理 ===")
+    
+    # 定义一个处理单个样本的函数
+    def process_sample(x):
+        # 一些复杂的变换
+        return jnp.sum(jnp.tanh(x) * jnp.cos(x))
+    
+    # 使用vmap进行向量化
+    batched_process = vmap(process_sample)
+    
+    # 创建批量数据
+    key = random.PRNGKey(3)
+    batch_size = 10000
+    sample_size = 1000
+    
+    batch_data = random.normal(key, (batch_size, sample_size))
+    
+    # 逐个处理（慢）
+    def sequential_process(data):
+        results = []
+        for i in range(data.shape[0]):
+            results.append(process_sample(data[i]))
+        return jnp.array(results)
+    
+    # 计时
+    start = time.time()
+    sequential_results = sequential_process(batch_data)
+    sequential_time = time.time() - start
+    
+    # 使用vmap批量处理（在GPU上并行）
+    start = time.time()
+    batched_results = batched_process(batch_data)
+    batched_time = time.time() - start
+    
+    print(f"批量大小: {batch_size}, 样本大小: {sample_size}")
+    print(f"逐个处理耗时: {sequential_time:.4f}秒")
+    print(f"批量处理耗时: {batched_time:.4f}秒")
+    print(f"加速比: {sequential_time/batched_time:.2f}倍")
+    
+    # 验证结果一致性
+    error = jnp.max(jnp.abs(sequential_results - batched_results))
+    print(f"结果最大误差: {error:.10f}")
+    
+    return batched_results
+
+# 5. 高级示例：简单的神经网络层
+def neural_network_example():
+    print("\n=== GPU上的神经网络层 ===")
+    
+    # 定义一个简单的神经网络层
+    def dense_layer(params, x):
+        w, b = params
+        return jnp.dot(x, w) + b
+    
+    # ReLU激活函数
+    def relu(x):
+        return jnp.maximum(0, x)
+    
+    # 两层神经网络
+    def two_layer_net(params, x):
+        w1, b1, w2, b2 = params
+        hidden = relu(jnp.dot(x, w1) + b1)
+        output = jnp.dot(hidden, w2) + b2
+        return output
+    
+    # 初始化参数
+    key = random.PRNGKey(4)
+    key1, key2, key3, key4 = random.split(key, 4)
+    
+    input_size = 784  # MNIST图像大小
+    hidden_size = 128
+    output_size = 10   # 10个类别
+    
+    # 初始化权重
+    w1 = random.normal(key1, (input_size, hidden_size)) * jnp.sqrt(2.0 / input_size)
+    b1 = jnp.zeros(hidden_size)
+    w2 = random.normal(key2, (hidden_size, output_size)) * jnp.sqrt(2.0 / hidden_size)
+    b2 = jnp.zeros(output_size)
+    
+    params = (w1, b1, w2, b2)
+    
+    # 创建批量输入
+    batch_size = 1024
+    x_batch = random.normal(key3, (batch_size, input_size))
+    
+    # JIT编译神经网络
+    jitted_net = jit(two_layer_net)
+    
+    # 前向传播
+    start = time.time()
+    outputs = jitted_net(params, x_batch)
+    inference_time = time.time() - start
+    
+    print(f"网络架构: {input_size} -> {hidden_size} -> {output_size}")
+    print(f"批量大小: {batch_size}")
+    print(f"GPU推理耗时: {inference_time:.6f}秒")
+    print(f"输出形状: {outputs.shape}")
+    print(f"每个样本的平均输出: {jnp.mean(outputs, axis=0)}")
+    
+    return outputs
+
+def main():
+    print("="*50)
+    print("JAX GPU Demo - 开始执行")
+    print("="*50)
+    
+    # 运行所有示例
+    try:
+        # 示例1: 基础GPU运算
+        result1 = basic_gpu_operations()
+        
+        # 示例2: JIT加速
+        result2 = jit_acceleration()
+        
+        # 示例3: 自动微分
+        result3 = autodiff_gpu()
+        
+        # 示例4: 批量处理
+        result4 = batch_processing()
+        
+        # 示例5: 神经网络
+        result5 = neural_network_example()
+        
+        print("\n" + "="*50)
+        print("所有示例执行完成!")
+        print("="*50)
+        
+        # 验证GPU确实在工作
+        print("\n验证GPU使用情况:")
+        with jax.profiler.trace("/tmp/jax-trace", create_perfetto_link=True):
+            # 执行一个操作来查看跟踪
+            test_matrix = jnp.ones((1000, 1000))
+            test_result = jnp.dot(test_matrix, test_matrix.T)
+            print(f"跟踪测试完成，结果形状: {test_result.shape}")
+            
+    except Exception as e:
+        print(f"执行出错: {e}")
+        print("\n注意：请确保已经正确安装JAX GPU版本")
+        print("安装命令: pip install --upgrade 'jax[cuda12_pip]' -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html")
+
+if __name__ == "__main__":
+    main()