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)

[3]:

[3]:

../_images/comparison_tianshou_batch.result_8_0.svg

[4]:

t.a

[4]:

tensor([[ 1.0407,  1.8661, -0.8391],
        [ 1.1429,  0.6467, -0.3856]])

[5]:

%timeit t.a

85.3 ns ± 0.621 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]:

../_images/comparison_tianshou_batch.result_11_0.svg

[7]:

%timeit t.a = new_value

91.3 ns ± 1.35 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]:

[9]:

Batch(
    a: tensor([[ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856]]),
    x: Batch(
           c: tensor([[ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665]]),
       ),
)

[10]:

b.a

[10]:

tensor([[ 1.0407,  1.8661, -0.8391],
        [ 1.1429,  0.6467, -0.3856]])

[11]:

%timeit b.a

78.5 ns ± 0.264 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.3942, -1.2681,  0.6362],
               [-2.2093,  1.3602,  1.6813]]),
    x: Batch(
           c: tensor([[ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665]]),
       ),
)

[13]:

%timeit b.a = new_value

641 ns ± 7.77 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

Initialization¶

TreeValue’s Initialization¶

[14]:

%timeit FastTreeValue(_TREE_DATA_1)

1.06 µs ± 5.25 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

Tianshou Batch’s Initialization¶

[15]:

%timeit Batch(**_TREE_DATA_1)

12.5 µs ± 47.6 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)

171 µs ± 3.61 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Deep Copy of Tianshou Batch¶

[18]:

b3 = Batch(**_TREE_DATA_3)
%timeit copy.deepcopy(b3)

173 µs ± 3.29 µs 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]:

../_images/comparison_tianshou_batch.result_34_0.svg

[21]:

%timeit t_stack(trees)

38.6 µs ± 366 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]:

../_images/comparison_tianshou_batch.result_36_0.svg

[23]:

%timeit t_cat(trees)

35.8 µs ± 298 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)

74.8 µs ± 275 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.2100, -0.1641, -1.6029,  0.6342],
                       [-0.0960, -0.2376,  0.2366, -0.6159],
                       [-0.6793, -0.9049, -0.1728,  0.2665]],

                      [[ 1.2100, -0.1641, -1.6029,  0.6342],
                       [-0.0960, -0.2376,  0.2366, -0.6159],
                       [-0.6793, -0.9049, -0.1728,  0.2665]],

                      [[ 1.2100, -0.1641, -1.6029,  0.6342],
                       [-0.0960, -0.2376,  0.2366, -0.6159],
                       [-0.6793, -0.9049, -0.1728,  0.2665]],

                      [[ 1.2100, -0.1641, -1.6029,  0.6342],
                       [-0.0960, -0.2376,  0.2366, -0.6159],
                       [-0.6793, -0.9049, -0.1728,  0.2665]],

                      [[ 1.2100, -0.1641, -1.6029,  0.6342],
                       [-0.0960, -0.2376,  0.2366, -0.6159],
                       [-0.6793, -0.9049, -0.1728,  0.2665]],

                      [[ 1.2100, -0.1641, -1.6029,  0.6342],
                       [-0.0960, -0.2376,  0.2366, -0.6159],
                       [-0.6793, -0.9049, -0.1728,  0.2665]],

                      [[ 1.2100, -0.1641, -1.6029,  0.6342],
                       [-0.0960, -0.2376,  0.2366, -0.6159],
                       [-0.6793, -0.9049, -0.1728,  0.2665]],

                      [[ 1.2100, -0.1641, -1.6029,  0.6342],
                       [-0.0960, -0.2376,  0.2366, -0.6159],
                       [-0.6793, -0.9049, -0.1728,  0.2665]]]),
       ),
    a: tensor([[[ 1.0407,  1.8661, -0.8391],
                [ 1.1429,  0.6467, -0.3856]],

               [[ 1.0407,  1.8661, -0.8391],
                [ 1.1429,  0.6467, -0.3856]],

               [[ 1.0407,  1.8661, -0.8391],
                [ 1.1429,  0.6467, -0.3856]],

               [[ 1.0407,  1.8661, -0.8391],
                [ 1.1429,  0.6467, -0.3856]],

               [[ 1.0407,  1.8661, -0.8391],
                [ 1.1429,  0.6467, -0.3856]],

               [[ 1.0407,  1.8661, -0.8391],
                [ 1.1429,  0.6467, -0.3856]],

               [[ 1.0407,  1.8661, -0.8391],
                [ 1.1429,  0.6467, -0.3856]],

               [[ 1.0407,  1.8661, -0.8391],
                [ 1.1429,  0.6467, -0.3856]]]),
)

[26]:

%timeit Batch.stack(batches)

98.9 µs ± 322 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

[27]:

Batch.cat(batches)

[27]:

Batch(
    x: Batch(
           c: tensor([[ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665],
                      [ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665],
                      [ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665],
                      [ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665],
                      [ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665],
                      [ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665],
                      [ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665],
                      [ 1.2100, -0.1641, -1.6029,  0.6342],
                      [-0.0960, -0.2376,  0.2366, -0.6159],
                      [-0.6793, -0.9049, -0.1728,  0.2665]]),
       ),
    a: tensor([[ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856],
               [ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856],
               [ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856],
               [ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856],
               [ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856],
               [ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856],
               [ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856],
               [ 1.0407,  1.8661, -0.8391],
               [ 1.1429,  0.6467, -0.3856]]),
)

[28]:

%timeit Batch.cat(batches)

180 µs ± 1.59 µ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))

442 µs ± 7.34 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[ ]: