在WSL2啟用cuda進行深度學習嘗鮮

在做深度學習相關實驗時，我一般會直接在本機 Linux 環境下進行，因為我需要基於 cuda 來啟用 GPU 運算。考慮到經常還要遷移到面向客戶PC機的智能服務，所以我偶爾也會選擇在 Windows 上進行。在 Windows 10 中自從有了 WSL，我的選擇是使用 WSL 遠程連接高性能服務器，這樣才能保持研發體驗的一致性，我特別期望能擺脫這樣的麻煩，如果能在Windows 10 本地切換一個 Linux系統又能使用 cuda 就好了，幸運的消息是，最新的 Windows 10 的更新推出 WSL2 對 cuda 的支持，我迫不及待地做了一下基准測試，想比較一下不同環境的性能表現。

我選擇的幾個系統環境是這樣：

• Ubuntu 20.04
• WSL2 Ubuntu 20.04
• Windows 10

主要硬件配置是這樣：

• Intel i9 10920X （12核，24線程）
• Nvidia RTX 2080 TI （11GB VRAM）

測試用的深度學習模型我使用的是最經典的MNIST，我比較熟悉Pytorch所以代碼我采用了這個示例倉庫，我修改了部分代碼如下，讓網絡模型更大以獲得更准確的讀數。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 128, 3, 1)
        self.conv2 = nn.Conv2d(128, 128, 3, 1)
        self.conv3 = nn.Conv2d(128, 128, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(15488, 15488//2)
        self.fc2 = nn.Linear(15488//2, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.conv3(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

其中batch size我設為512，epochs設為14，以precision為FP32運行，結果對比如下：

  

哈哈，老實說結果還湊合！WSL2 在啟用 cuda 支持后比本機 Ubuntu 環境下多消耗了18%，要知道這是在 Nvidia Rtx 2080 Ti 上訓練模型。目前 WSL2 對 cuda 的支持仍處於早期預覽模式，期待各路工程師、研究員大神們，以及微軟和 Nvidia 再打磨下，盡快達到接近本機 Ubuntu 性能的水平。如果對訓練時長消耗特別嚴格的研究（比如海量圖像、語料）來說，多花18%是很難接受的，但對我來說，訓練樣本數據都還比較小，一般也就到2G差不多了，所以我可以稍稍犧牲下操作系統層面的性能，采用異步訓練來彌補一下，畢竟我換來的是能用 Windows 10 和 WSL2 作為我的主研發環境，不用多系統或者遠程Linux系統來回切換了。