Comparison Between TreeValue and Tianshou Batch¶
In this section, we will take a look at the feature and performance of the Tianshou Batch library, which is developed by Tsinghua Machine Learning Group.
Before starting the comparison, let us define some thing.
[1]:
import torch
_TREE_DATA_1 = {'a': 1, 'b': 2, 'x': {'c': 3, 'd': 4}}
_TREE_DATA_2 = {
'a': torch.randn(2, 3),
'x': {
'c': torch.randn(3, 4)
},
}
_TREE_DATA_3 = {
'obs': torch.randn(4, 84, 84),
'action': torch.randint(0, 6, size=(1,)),
'reward': torch.rand(1),
}
Read and Write Operation¶
Reading and writing are the two most common operations in the tree data structure based on the data model (TreeValue and Tianshou Batch both belong to this type), so this section will compare the reading and writing performance of these two libraries.
TreeValue’s Get and Set¶
[2]:
from treevalue import FastTreeValue
t = FastTreeValue(_TREE_DATA_2)
/tmp/tmp4q2ykbbk/fbb7ad1854f8455286e2cedd40b3312212c7f17a/treevalue/tree/integration/torch.py:18: FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.
register_for_torch(TreeValue)
/tmp/tmp4q2ykbbk/fbb7ad1854f8455286e2cedd40b3312212c7f17a/treevalue/tree/integration/torch.py:19: FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.
register_for_torch(FastTreeValue)
[3]:
t
[3]:
<FastTreeValue 0x7f54f09764c0>
├── 'a' --> tensor([[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]])
└── 'x' --> <FastTreeValue 0x7f54f0976f70>
└── 'c' --> tensor([[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]])
[4]:
t.a
[4]:
tensor([[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]])
[5]:
%timeit t.a
48.7 ns ± 0.673 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
[6]:
new_value = torch.randn(2, 3)
t.a = new_value
t
[6]:
<FastTreeValue 0x7f54f09764c0>
├── 'a' --> tensor([[-0.1345, -1.2116, 1.2098],
│ [ 1.2173, -1.8312, -0.2737]])
└── 'x' --> <FastTreeValue 0x7f54f0976f70>
└── 'c' --> tensor([[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]])
[7]:
%timeit t.a = new_value
56.7 ns ± 0.0667 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
Tianshou Batch’s Get and Set¶
[8]:
from tianshou.data import Batch
b = Batch(**_TREE_DATA_2)
[9]:
b
[9]:
Batch(
a: tensor([[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]]),
x: Batch(
c: tensor([[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]]),
),
)
[10]:
b.a
[10]:
tensor([[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]])
[11]:
%timeit b.a
41 ns ± 0.362 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
[12]:
new_value = torch.randn(2, 3)
b.a = new_value
b
[12]:
Batch(
a: tensor([[-1.8601, 2.1694, -1.1358],
[ 0.0145, -1.6993, 0.2390]]),
x: Batch(
c: tensor([[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]]),
),
)
[13]:
%timeit b.a = new_value
367 ns ± 0.489 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
Initialization¶
TreeValue’s Initialization¶
[14]:
%timeit FastTreeValue(_TREE_DATA_1)
625 ns ± 6.03 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
Tianshou Batch’s Initialization¶
[15]:
%timeit Batch(**_TREE_DATA_1)
8.68 µs ± 15.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
Deep Copy Operation¶
[16]:
import copy
Deep Copy of TreeValue¶
[17]:
t3 = FastTreeValue(_TREE_DATA_3)
%timeit copy.deepcopy(t3)
128 µs ± 495 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
Deep Copy of Tianshou Batch¶
[18]:
b3 = Batch(**_TREE_DATA_3)
%timeit copy.deepcopy(b3)
127 µs ± 526 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
Stack, Concat and Split Operation¶
Performance of TreeValue¶
[19]:
trees = [FastTreeValue(_TREE_DATA_2) for _ in range(8)]
[20]:
t_stack = FastTreeValue.func(subside=True)(torch.stack)
t_stack(trees)
[20]:
<FastTreeValue 0x7f54f09dd6d0>
├── 'a' --> tensor([[[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]],
│
│ [[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]],
│
│ [[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]],
│
│ [[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]],
│
│ [[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]],
│
│ [[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]],
│
│ [[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]],
│
│ [[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]]])
└── 'x' --> <FastTreeValue 0x7f54f09c5f40>
└── 'c' --> tensor([[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]]])
[21]:
%timeit t_stack(trees)
23.8 µs ± 231 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
[22]:
t_cat = FastTreeValue.func(subside=True)(torch.cat)
t_cat(trees)
[22]:
<FastTreeValue 0x7f54f09dd760>
├── 'a' --> tensor([[ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482],
│ [ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482],
│ [ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482],
│ [ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482],
│ [ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482],
│ [ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482],
│ [ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482],
│ [ 0.5140, 2.0014, 0.4726],
│ [-0.3070, 1.5447, 0.5482]])
└── 'x' --> <FastTreeValue 0x7f53f9e2eb20>
└── 'c' --> tensor([[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]])
[23]:
%timeit t_cat(trees)
22.2 µs ± 612 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
[24]:
t_split = FastTreeValue.func(rise=True)(torch.split)
tree = FastTreeValue({
'obs': torch.randn(8, 4, 84, 84),
'action': torch.randint(0, 6, size=(8, 1,)),
'reward': torch.rand(8, 1),
})
%timeit t_split(tree, 1)
50.5 µs ± 383 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
Performance of Tianshou Batch¶
[25]:
batches = [Batch(**_TREE_DATA_2) for _ in range(8)]
Batch.stack(batches)
[25]:
Batch(
x: Batch(
c: tensor([[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]],
[[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]]]),
),
a: tensor([[[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]],
[[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]],
[[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]],
[[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]],
[[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]],
[[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]],
[[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]],
[[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]]]),
)
[26]:
%timeit Batch.stack(batches)
63.6 µs ± 924 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
[27]:
Batch.cat(batches)
[27]:
Batch(
x: Batch(
c: tensor([[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524],
[ 1.8188, 1.8288, 0.6108, 0.7464],
[-0.8200, -0.1059, 0.4295, 1.1611],
[ 0.0272, -0.6516, -0.2357, 0.9524]]),
),
a: tensor([[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482],
[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482],
[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482],
[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482],
[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482],
[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482],
[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482],
[ 0.5140, 2.0014, 0.4726],
[-0.3070, 1.5447, 0.5482]]),
)
[28]:
%timeit Batch.cat(batches)
117 µs ± 1.05 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
[29]:
batch = Batch({
'obs': torch.randn(8, 4, 84, 84),
'action': torch.randint(0, 6, size=(8, 1,)),
'reward': torch.rand(8, 1)}
)
%timeit list(Batch.split(batch, 1, shuffle=False, merge_last=True))
280 µs ± 2.08 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
[ ]: