thu-ml · Trinkle23897 · Aug 2, 2020 · Jul 11, 2020 · Jul 11, 2020 · Jul 12, 2020
diff --git a/.github/workflows/profile.yml b/.github/workflows/profile.yml
@@ -0,0 +1,22 @@
+name: Data Profile
+
+on: [push, pull_request]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+    - name: Upgrade pip
+      run: |
+        python -m pip install --upgrade pip setuptools wheel
+    - name: Install dependencies
+      run: |
+        pip install ".[dev]" --upgrade
+    - name: Test with pytest
+      run: |
+        pytest test/throughput --durations=0 -v
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
@@ -28,8 +28,9 @@ jobs:
       run: |
         pip install ".[dev]" --upgrade
     - name: Test with pytest
+      # ignore test/throughput which only profiles the code
       run: |
-        pytest test --cov tianshou --cov-report=xml --durations 0 -v
+        pytest test --ignore-glob='*profile.py' --cov=tianshou --cov-report=xml --durations=0 -v
     - name: Upload coverage to Codecov
       uses: codecov/codecov-action@v1
       with:

diff --git a/test/throughput/__init__.py b/test/throughput/__init__.py
diff --git a/test/throughput/test_batch_profile.py b/test/throughput/test_batch_profile.py
@@ -0,0 +1,120 @@
+import copy
+import pickle
+
+import numpy as np
+import pytest
+import torch
+
+from tianshou.data import Batch
+
+
+@pytest.fixture(scope="module")
+def data():
+    print("Initialising data...")
+    np.random.seed(0)
+    batch_set = [Batch(a=[j for j in np.arange(1e3)],
+                       b={'b1': (3.14, 3.14), 'b2': np.arange(1e3)},
+                       c=i) for i in np.arange(int(1e4))]
+    batch0 = Batch(
+        a=np.ones((3, 4), dtype=np.float64),
+        b=Batch(
+            c=np.ones((1,), dtype=np.float64),
+            d=torch.ones((3, 3, 3), dtype=torch.float32),
+            e=list(range(3))
+        )
+    )
+    batchs1 = [copy.deepcopy(batch0) for _ in np.arange(1e4)]
+    batchs2 = [copy.deepcopy(batch0) for _ in np.arange(1e4)]
+    batch_len = int(1e4)
+    batch3 = Batch(obs=[np.arange(20) for _ in np.arange(batch_len)],
+                   reward=np.arange(batch_len))
+    indexs = np.random.choice(batch_len,
+                              size=batch_len//10, replace=False)
+    slice_dict = {'obs': [np.arange(20)
+                          for _ in np.arange(batch_len//10)],
+                  'reward': np.arange(batch_len//10)}
+    dict_set = [{'obs': np.arange(20), 'info': "this is info", 'reward': 0}
+                for _ in np.arange(1e2)]
+    batch4 = Batch(
+        a=np.ones((10000, 4), dtype=np.float64),
+        b=Batch(
+            c=np.ones((1,), dtype=np.float64),
+            d=torch.ones((1000, 1000), dtype=torch.float32),
+            e=np.arange(1000)
+        )
+    )
+
+    print("Initialised")
+    return {'batch_set': batch_set,
+            'batch0': batch0,
+            'batchs1': batchs1,
+            'batchs2': batchs2,
+            'batch3': batch3,
+            'indexs': indexs,
+            'dict_set': dict_set,
+            'slice_dict': slice_dict,
+            'batch4': batch4
+            }
+
+
+def test_init(data):
+    """Test Batch __init__()."""
+    for _ in np.arange(10):
+        _ = Batch(data['batch_set'])
+
+
+def test_get_item(data):
+    """Test get with item."""
+    for _ in np.arange(1e5):
+        _ = data['batch3'][data['indexs']]
+
+
+def test_get_attr(data):
+    """Test get with attr."""
+    for _ in np.arange(1e6):
+        data['batch3'].get('obs')
+        data['batch3'].get('reward')
+        _, _ = data['batch3'].obs, data['batch3'].reward
+
+
+def test_set_item(data):
+    """Test set with item."""
+    for _ in np.arange(1e4):
+        data['batch3'][data['indexs']] = data['slice_dict']
+
+
+def test_set_attr(data):
+    """Test set with attr."""
+    for _ in np.arange(1e4):
+        data['batch3'].c = np.arange(1e3)
+        data['batch3'].obs = data['dict_set']
+
+
+def test_numpy_torch_convert(data):
+    """Test conversion between numpy and torch."""
+    for _ in np.arange(1e5):
+        data['batch4'].to_torch()
+        data['batch4'].to_numpy()
+
+
+def test_pickle(data):
+    for _ in np.arange(1e4):
+        pickle.loads(pickle.dumps(data['batch4']))
+
+
+def test_cat(data):
+    """Test cat"""
+    for i in range(10000):
+        Batch.cat((data['batch0'], data['batch0']))
+        data['batchs1'][i].cat_(data['batch0'])
+
+
+def test_stack(data):
+    """Test stack"""
+    for i in range(10000):
+        Batch.stack((data['batch0'], data['batch0']))
+        data['batchs2'][i].stack_([data['batch0']])
+
+
+if __name__ == '__main__':
+    pytest.main(["-s", "-k batch_profile", "--durations=0", "-v"])
diff --git a/test/throughput/test_buffer_profile.py b/test/throughput/test_buffer_profile.py
@@ -0,0 +1,81 @@
+import numpy as np
+import pytest
+
+from tianshou.data import (ListReplayBuffer, PrioritizedReplayBuffer,
+                           ReplayBuffer)
+
+
+@pytest.fixture(scope="module")
+def data():
+    np.random.seed(0)
+    obs = {'observable': np.random.rand(
+        100, 100), 'hidden': np.random.randint(1000, size=200)}
+    info = {'policy': "dqn", 'base': np.arange(10)}
+    add_data = {'obs': obs, 'rew': 1., 'act': np.random.rand(30),
+                'done': False, 'obs_next': obs, 'info': info}
+    buffer = ReplayBuffer(int(1e3), stack_num=100)
+    buffer2 = ReplayBuffer(int(1e4), stack_num=100)
+    indexes = np.random.choice(int(1e3), size=3, replace=False)
+    return{
+        'add_data': add_data,
+        'buffer': buffer,
+        'buffer2': buffer2,
+        'slice': slice(-3000, -1000, 2),
+        'indexes': indexes
+    }
+
+
+def test_init():
+    for _ in np.arange(1e5):
+        _ = ReplayBuffer(1e5)
+        _ = PrioritizedReplayBuffer(
+            size=int(1e5), alpha=0.5,
+            beta=0.5, repeat_sample=True)
+        _ = ListReplayBuffer()
+
+
+def test_add(data):
+    buffer = data['buffer']
+    for _ in np.arange(1e5):
+        buffer.add(**data['add_data'])
+
+
+def test_update(data):
+    buffer = data['buffer']
+    buffer2 = data['buffer2']
+    for _ in np.arange(1e2):
+        buffer2.update(buffer)
+
+
+def test_getitem_slice(data):
+    Slice = data['slice']
+    buffer = data['buffer']
+    for _ in np.arange(1e3):
+        _ = buffer[Slice]
+
+
+def test_getitem_indexes(data):
+    indexes = data['indexes']
+    buffer = data['buffer']
+    for _ in np.arange(1e2):
+        _ = buffer[indexes]
+
+
+def test_get(data):
+    indexes = data['indexes']
+    buffer = data['buffer']
+    for _ in np.arange(3e2):
+        buffer.get(indexes, 'obs')
+        buffer.get(indexes, 'rew')
+        buffer.get(indexes, 'done')
+        buffer.get(indexes, 'info')
+
+
+def test_sample(data):
+    buffer = data['buffer']
+    for _ in np.arange(1e1):
+        buffer.sample(int(1e2))
+
+
+if __name__ == '__main__':
+    pytest.main(["-s", "-k buffer_profile", "--durations=0", "-v"])