PyTorch 的 C++ API LibTorch 中是否有对应的 torch.distributions.Normal?
pytorch 444
原文标题 :Is there an equivalent of torch.distributions.Normal in LibTorch, the C++ API for PyTorch?
我正在使用随机策略实现策略梯度算法,并且由于 Python 中的“辅助”非 PyTorch 操作很慢,我想在 C++ 中实现该算法。有没有办法在 PyTorch C++ API 中实现正态分布?
回复
我来回复-
vaporK 评论
Python 实现实际上调用了
at::
命名空间中的 C++ 后端(CPU、CUDA,我在其中找到了这个)。在 PyTorch 团队和/或贡献者在 LibTorch 中实现前端之前,您可以使用类似这样的方法来解决它(我只实现了rsample()
和log_prob()
,因为这是我在这个用例中需要的):constexpr double lz = log(sqrt(2 * M_PI)); class Normal { torch::Tensor mean, stddev, var, log_std; public: Normal(const torch::Tensor &mean, const torch::Tensor &std) : mean(mean), stddev(std), var(std * std), log_std(std.log()) {} torch::Tensor rsample() { auto device = torch::cuda::is_available() ? torch::kCUDA : torch::kCPU; auto eps = torch::randn(1).to(device); return this->mean + eps * this->stddev; } torch::Tensor log_prob(const torch::Tensor &value) { // log [exp(-(x-mu)^2/(2 sigma^2)) / (sqrt(2 pi) * sigma)] = // = log [exp(-(x-mu)^2/(2 sigma^2))] - log [sqrt(2 pi) * sigma] = // = -(x - mu)^2 / (2 sigma^2) - log(sigma) - log(sqrt(2 pi)) return -(value - this->mean)*(value - this->mean) / (2 * this->var) - this->log_std - lz; } };
2年前