treequeues: 為pytree物件提供高效能的佇列

banq發表於2022-04-24

如果您使用 jax 並且需要在程式之間傳遞一些 pytree,我可能會為您提供一些東西:)
我開發了一個“樹佇列”。它是為 pytree 的巢狀陣列建立的佇列。
傳輸速度比普通佇列快10倍。這是透過利用共享記憶體陣列和避免pickling資料來完成的。這在開發分散式架構時非常有用,例如,速度是最重要的分散式強化學習。
在我的例子中,這個實現對於在實現 PBT 演算法時消除瓶頸非常有用!

這個庫包包含了使用pytree和multiprocessing.Arrays在程式之間傳輸陣列和巢狀陣列的佇列。與vanilla multiprocessing.Queue相比,這個實現可以達到高達10倍的速度,這取決於樹的形狀和大小以及涉及的程式數量。

透過使用numpy陣列與多程式陣列的緩衝,資料可以在不需要醃製的情況下傳送。
缺點之一是總大小(巢狀陣列的大小與佇列的最大大小)需要預先分配。

這個包包含TreeQueue和ArrayQueue,在這兩種情況下,建立佇列時需要傳遞一個資料例項和最大尺寸。

這個資源庫的靈感來自portugueslab的ArrayQueues。

點選標題

相關文章