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()