Windows10下PyTorch调用torch.compile报错：暂不支持Windows系统

背景 #

时隔23年重返校园攻读博士学位，本学期修习深度学习课程。1999-2001年研究生阶段曾用C语言大量开发神经网络，当前正学习Python及相关深度学习环境。此前尝试配置TensorFlow（需WSL2运行），耗时约20小时未成功，转而使用PyTorch，但无论用Anaconda安装、PyTorch官方Anaconda包还是pip安装，均遭遇相同错误：RuntimeError: Windows not yet supported for torch.compile。

环境信息 #

NVIDIA显卡信息 #

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 546.12                 Driver Version: 546.12       CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+| GPU  Name                     TCC/WDDM  | Bus-Id        Disp.A | Volatile Uncorr. ECC || Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. ||                                         |                      |               MIG M. ||=========================================+======================+======================||   0  Quadro M4000M                WDDM  | 00000000:01:00.0 Off |                  N/A || N/A   42C    P0              26W / 100W |      0MiB /  4096MiB |      0%      Default ||                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            ||  GPU   GI   CI        PID   Type   Process name                            GPU Memory ||        ID   ID                                                             Usage      ||=======================================================================================||  No running processes found                                                           |
+---------------------------------------------------------------------------------------+

安装脚本（新.venv环境） #

pip install jupyter
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install pandas
pip install scikit-learn

环境检测代码及结果 #

检测代码：

import sys
import platform
import torch
import pandas as pd
import sklearn as sk

has_gpu = torch.cuda.is_available()
has_mps = getattr(torch,'has_mps',False)
device = "mps" if getattr(torch,'has_mps',False) else "gpu" if torch.cuda.is_available() else "cpu"

print(f"Python Platform: {platform.platform()}")
print(f"PyTorch Version: {torch.__version__}")
print()
print(f"Python {sys.version}")
print(f"system platform: {sys.platform} {platform.architecture()}")
print(f"Pandas {pd.__version__}")
print(f"Scikit-Learn {sk.__version__}")
print("NVIDIA/CUDA GPU is", "available" if has_gpu else "NOT AVAILABLE")    
print(f"Target device is {device}")

检测结果：

Python Platform: Windows-10-10.0.19045-SP0
PyTorch Version: 2.1.2+cu121

Python 3.11.6 (tags/v3.11.6:8b6ee5b, Oct  2 2023, 14:57:12) [MSC v.1935 64 bit (AMD64)]
system platform: win32 ('64bit', 'WindowsPE')
Pandas 2.1.4
Scikit-Learn 1.3.2
NVIDIA/CUDA GPU is available
Target device is gpu

问题复现 #

运行以下测试代码时触发错误：

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(100, 10)

    def forward(self, x):
        return torch.nn.functional.relu(self.lin(x))

mod = MyModule()
opt_mod = torch.compile(mod)
print(opt_mod(torch.randn(10, 100)))

错误信息：

RuntimeError: Windows not yet supported for torch.compile

已在3台电脑、Python 3.8/3.9/3.10/3.12版本中复现该问题，查看PyTorch源码发现明确限制：

def check_if_dynamo_supported():
    if sys.platform == "win32":
        raise RuntimeError("Windows not yet supported for torch.compile")
    if sys.version_info >= (3, 12):
        raise RuntimeError("Python 3.12+ not yet supported for torch.compile")

解决方案 #

1. 直接弃用torch.compile：PyTorch当前版本确实未在Windows平台支持torch.compile，但CUDA功能正常可用。只需移除opt_mod = torch.compile(mod)代码，直接调用mod(torch.randn(10, 100))即可正常运行模型，不影响训练和推理。
2. 切换至WSL2环境：若必须使用torch.compile，可安装WSL2并在Linux子系统中配置PyTorch，既能解决torch.compile的支持问题，也能解决之前TensorFlow的配置难题。
3. 等待官方更新：PyTorch团队大概率会在后续版本中添加Windows对torch.compile的支持，可关注官方版本更新日志。

内容的提问来源于stack exchange，提问作者Kris Jensen