mirror of
https://github.com/wassname/ray.git
synced 2026-07-30 12:30:30 +08:00
Siyuan (Ryans) ZhuangandGitHub
3f22448834
* [Core] zero-copy serializer for pytorch (#12344)
* zero-copy serializer for pytorch
* address possible bottleneck
* add tests & device support
(cherry picked from commit 0a505ca83d)
* add environmental variables
* update doc
30 lines
787 B
Python
Executable File
30 lines
787 B
Python
Executable File
#!/usr/bin/env python
|
|
|
|
import os
|
|
import sys
|
|
|
|
if __name__ == "__main__":
|
|
# Do not import torch for testing purposes.
|
|
os.environ["RLLIB_TEST_NO_TORCH_IMPORT"] = "1"
|
|
os.environ["RAY_DISABLE_PYTORCH_SERIALIZER"] = "1"
|
|
|
|
from ray.rllib.agents.a3c import A2CTrainer
|
|
assert "torch" not in sys.modules, \
|
|
"Torch initially present, when it shouldn't."
|
|
|
|
# note: no ray.init(), to test it works without Ray
|
|
trainer = A2CTrainer(
|
|
env="CartPole-v0", config={
|
|
"framework": "tf",
|
|
"num_workers": 0
|
|
})
|
|
trainer.train()
|
|
|
|
assert "torch" not in sys.modules, "Torch should not be imported"
|
|
|
|
# Clean up.
|
|
del os.environ["RLLIB_TEST_NO_TORCH_IMPORT"]
|
|
del os.environ["RAY_DISABLE_PYTORCH_SERIALIZER"]
|
|
|
|
print("ok")
|