Skip to content

[Feature]: faster safetanh transform via C++ bindings #289

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 18, 2022
Merged

[Feature]: faster safetanh transform via C++ bindings #289

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a91c591
init
vmoens Jul 18, 2022
4d8e118
init
vmoens Jul 18, 2022
af5d2c5
amend
vmoens Jul 18, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion test/test_distributions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import pytest
import torch
from _utils_internal import get_available_devices
from torch import nn
from torch import nn, autograd
from torchrl.data.tensordict.tensordict import _TensorDict
from torchrl.modules import (
TanhNormal,
Expand Down Expand Up @@ -211,6 +211,18 @@ def test_tanhtrsf(dtype):
assert (some_big_number.sign() == ones.sign()).all()


@pytest.mark.parametrize("dtype", [torch.float, torch.double])
def test_tanhtrsf_grad(dtype):
torch.manual_seed(0)
trsf = SafeTanhTransform()
x = torch.randn(100, requires_grad=True)
y1 = trsf(x)
y2 = x.tanh()
g1 = autograd.grad(y1.sum(), x, retain_graph=True)[0]
g2 = autograd.grad(y2.sum(), x, retain_graph=True)[0]
torch.testing.assert_close(g1, g2)


if __name__ == "__main__":
args, unknown = argparse.ArgumentParser().parse_known_args()
pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)
3 changes: 3 additions & 0 deletions torchrl/csrc/pybind.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include <memory>

#include "segment_tree.h"
#include "utils.h"

namespace py = pybind11;

Expand All @@ -20,4 +21,6 @@ PYBIND11_MODULE(_torchrl, m) {

torchrl::DefineMinSegmentTree<float>("Fp32", m);
torchrl::DefineMinSegmentTree<double>("Fp64", m);

m.def("safetanh", &safetanh, "Safe Tanh");
}
13 changes: 13 additions & 0 deletions torchrl/csrc/utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

#include <torch/extension.h>
#include <torch/torch.h>

#include <iostream>

torch::Tensor safetanh(torch::Tensor input) {
return torch::clamp(torch::tanh(input), -0.999999, 0.999999);
}
5 changes: 2 additions & 3 deletions torchrl/modules/distributions/continuous.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from torch import distributions as D, nn
from torch.distributions import constraints

from torchrl._torchrl import safetanh
from torchrl.modules.distributions.truncated_normal import (
TruncatedNormal as _TruncatedNormal,
)
Expand Down Expand Up @@ -85,9 +86,7 @@ class SafeTanhTransform(D.TanhTransform):
"""

def _call(self, x: torch.Tensor) -> torch.Tensor:
eps = torch.finfo(x.dtype).eps
y = super()._call(x)
y.data.clamp_(-1 + eps, 1 - eps)
y = safetanh(x)
return y

def _inverse(self, y: torch.Tensor) -> torch.Tensor:
Expand Down