1. 硬盘 <-> CPU 数据传输速度

import time
import os

# 定义文件大小和测试文件路径
file_size = 1 * 1024 * 1024 * 100  # 100 MB 的文件大小
file_path = "test_file.bin"

# 创建一个测试文件并测量写入速度
def test_write_speed():
    data = os.urandom(file_size)  # 生成随机数据
    start_time = time.time()  # 记录开始时间

    with open(file_path, 'wb') as f:
        f.write(data)

    end_time = time.time()  # 记录结束时间
    write_time = end_time - start_time
    write_speed = file_size / (write_time * 1024 * 1024)  # 转换为 MB/s
    print(f"Disk write speed: {write_speed:.6f} MB/s")

# 测试硬盘读取速度
def test_read_speed():
    start_time = time.time()  # 记录开始时间

    with open(file_path, 'rb') as f:
        data = f.read()

    end_time = time.time()  # 记录结束时间
    read_time = end_time - start_time
    read_speed = file_size / (read_time * 1024 * 1024)  # 转换为 MB/s
    print(f"Disk read speed: {read_speed:.6f} MB/s")

# 进行测试
test_write_speed()
test_read_speed()

# 删除测试文件
os.remove(file_path)

2. CPU <-> GPU 数据传输速度

import torch
import time

# 设置设备
device_cpu = torch.device('cpu')
device_gpu = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 创建一个随机张量在 CPU 上
size = 10**7  # 可以根据需要调整大小，元素数量
data_cpu = torch.randn(size).to(device_cpu)

# 计算数据大小，假设是 float32 类型，每个元素 4 字节
data_size_bytes = data_cpu.nelement() * data_cpu.element_size()

# 测量 CPU -> GPU 传输时间
start_time = time.time()
data_gpu = data_cpu.to(device_gpu)
torch.cuda.synchronize()  # 确保 GPU 操作完成
cpu_to_gpu_time = time.time() - start_time

# 计算传输速度 (MB/s)
cpu_to_gpu_speed = data_size_bytes / (cpu_to_gpu_time * 1024 * 1024)  # 转换为 MB/s
print(f"CPU -> GPU data transfer speed: {cpu_to_gpu_speed:.6f} MB/s")

# 测量 GPU -> CPU 传输时间
start_time = time.time()
data_back_to_cpu = data_gpu.to(device_cpu)
torch.cuda.synchronize()  # 确保 GPU 操作完成
gpu_to_cpu_time = time.time() - start_time

# 计算传输速度 (MB/s)
gpu_to_cpu_speed = data_size_bytes / (gpu_to_cpu_time * 1024 * 1024)  # 转换为 MB/s
print(f"GPU -> CPU data transfer speed: {gpu_to_cpu_speed:.6f} MB/s")