download_size_torch.py

import os
import tempfile
import subprocess
import matplotlib.pyplot as plt

# List of PyTorch versions to download
torch_versions = [
    "torch==1.11.0",
    "torch==1.12.0",
    "torch==1.12.1",
    "torch==1.13.0",
    "torch==1.13.1",
    "torch==2.0.0",
]

def download_torch_version(version):
    with tempfile.TemporaryDirectory() as temp_dir:
        subprocess.run(["pip", "download", version, "-d", temp_dir])
        total_size = 0
        for file in os.listdir(temp_dir):
            file_path = os.path.join(temp_dir, file)
            total_size += os.path.getsize(file_path)
        return total_size

sizes = [download_torch_version(version) for version in torch_versions]

# Convert sizes to MB
sizes_mb = [size / (1024 * 1024) for size in sizes]

# Plot the sizes
plt.plot(torch_versions, sizes_mb, marker="o")
plt.xlabel("PyTorch Version")
plt.ylabel("Size (MB)")
plt.title("PyTorch Version vs. Size")
plt.xticks(rotation=45)
plt.grid()
plt.show()
plt.savefig("torch_version_vs_size.png")