在固定某些層的多個 GPU 上訓練單個 pytorch 模型？

Question

我在使用 pytorch DistributedDataParallel時遇到了一些問題。 情況是：

1. 我的模型是A ，它像往常一樣在單個GPU 上訓練。 假設有三層答：

    class A(nn.module): def __init__(self): super(A,self).__init__() self.layer0 = layer0 self.layer1 = layer1 self.layer2 = layer2 def forward(self,x): x=self.layer0(x) x=self.layer1(x) x=self.layer2(x) return x

2. 現在我有了一些新數據。 我想在多個GPU 上用它微調A。 我需要將A包裝為多 GPU 模型B 。

3. 但是有兩個訓練階段。 在第一階段，我想修復B 的layer0和layer1 。 在第二階段，只修復layer0 。 然后requires_grad在訓練期間更改layer1 requires_grad參數。 但是， DistributedDataParallel文檔說：

   在使用 DistributedDataParallel 包裝模型后，您永遠不應該嘗試更改模型的參數。

事實上，我試圖用B.module指一個包裹在B中。 但與單GPU模型相比，測試結果異常。 也許這種方式是不允許的。

我該怎么辦？ 有什么合適的方法來包裝我的模型嗎？ 保存和加載模型時應該注意什么？

只需在具有多個 GPU 的單台機器上運行它，這樣您就可以忽略使用多台機器的分布式情況。 非常感謝。

2019.12.03 更新

正如@jodag 所建議的，我嘗試了DataParallel ，但沒有用。 這次我在包裝之后沒有改變B中的任何東西（除了訓練它）。 為簡單起見，我的代碼是這樣的（我參考了這個）：

class B(nn.DataParallel):
     def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.module, name)
a = A()
b = B(a,device_ids=[0,1])
b = b.cuda()
trained_param = b.layer2.parameters()
# trained_param = [{'params':b.layer2.parameters()},{'params':b.layer1.parameters()}]
optimizer = optim.Adam(trained_param)
b.train()
...
for x, label in data_loader:
    optimizer.zero_grad()
    x = x.to(0) # This line can be commented.
    y = b(x)
    l = loss(y, label)
    l.backword()
    optimizer.step()

Answer 1

如果您只嘗試優化部分參數，為什么不嘗試通過優化器而不是模型來控制它？
您可以保持模型原樣（包裝在DistributedDataParallel ）並僅將其部分參數傳遞給相關優化器。