diff --git a/src/nn/activations.rs b/src/nn/activations.rs index 8a9e4860..853abccf 100644 --- a/src/nn/activations.rs +++ b/src/nn/activations.rs @@ -36,6 +36,7 @@ activation_impls!(Cos, try_cos, #[doc="Calls [cos()]."]); activation_impls!(Ln, try_ln, #[doc="Calls [ln()]."]); activation_impls!(Exp, try_exp, #[doc="Calls [exp()]."]); activation_impls!(Sigmoid, try_sigmoid, #[doc="Calls [sigmoid()]."]); +activation_impls!(Sigmoidr, try_sigmoidr, #[doc="Calls [sigmoidr()]."]); activation_impls!(Tanh, try_tanh, #[doc="Calls [tanh()]."]); activation_impls!(Square, try_square, #[doc="Calls [square()]."]); activation_impls!(Sqrt, try_sqrt, #[doc="Calls [sqrt()]."]); diff --git a/src/tensor_ops/mod.rs b/src/tensor_ops/mod.rs index 0cd5cc90..8bd13978 100644 --- a/src/tensor_ops/mod.rs +++ b/src/tensor_ops/mod.rs @@ -195,6 +195,7 @@ mod roll; mod select_and_gather; mod sgd; mod sigmoid; +mod sigmoidr; mod sin; mod slice; mod softmax; @@ -259,6 +260,7 @@ pub use roll::Roll; pub use select_and_gather::{GatherTo, SelectTo}; pub use sgd::SgdConfig; pub use sigmoid::sigmoid; +pub use sigmoidr::sigmoidr; pub use sin::sin; pub use slice::slice; pub use softmax::softmax; diff --git a/src/tensor_ops/sigmoidr/cpu_kernel.rs b/src/tensor_ops/sigmoidr/cpu_kernel.rs new file mode 100644 index 00000000..a1de352a --- /dev/null +++ b/src/tensor_ops/sigmoidr/cpu_kernel.rs @@ -0,0 +1,15 @@ +use crate::tensor_ops::cpu_kernels::UnaryDerivative; + +impl UnaryDerivative for super::SigmoidrKernelOp { + const DF_USES_FX: bool = true; + const HAS_CONST_DF: bool = false; + #[inline(always)] + fn f(&self, x: &F) -> F { + F::one() / (F::one() + x.neg().exp()) + } + #[inline(always)] + fn df(&self, &fx: &F) -> F { + let d = fx * (F::one() - fx); + F::max(d, F::from(0.0000001).unwrap()) + } +} diff --git a/src/tensor_ops/sigmoidr/cuda_kernel.rs b/src/tensor_ops/sigmoidr/cuda_kernel.rs new file mode 100644 index 00000000..fc360a9b --- /dev/null +++ b/src/tensor_ops/sigmoidr/cuda_kernel.rs @@ -0,0 +1,15 @@ +use super::SigmoidrKernelOp as Sigmoidr; +#[allow(unused_imports)] +use crate::dtypes::*; +use crate::tensor_ops::cuda_kernels::cuda_unary; + +unsafe impl cudarc::driver::DeviceRepr for Sigmoidr {} + +const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/sigmoidr.ptx")); + +#[cfg(feature = "f16")] +cuda_unary!(df(f(x)) Sigmoidr, f16, PTX, "sigmoidr_fwd_f16", "sigmoidr_bwd_f16"); +#[cfg(feature = "f16")] +cuda_unary!(df(f(x)) Sigmoidr, AMP, PTX, "sigmoidr_fwd_f16", "sigmoidr_bwd_f16"); +cuda_unary!(df(f(x)) Sigmoidr, f32, PTX, "sigmoidr_fwd_f32", "sigmoidr_bwd_f32"); +cuda_unary!(df(f(x)) Sigmoidr, f64, PTX, "sigmoidr_fwd_f64", "sigmoidr_bwd_f64"); diff --git a/src/tensor_ops/sigmoidr/mod.rs b/src/tensor_ops/sigmoidr/mod.rs new file mode 100644 index 00000000..146335e6 --- /dev/null +++ b/src/tensor_ops/sigmoidr/mod.rs @@ -0,0 +1,59 @@ +mod cpu_kernel; + +#[cfg(feature = "cuda")] +mod cuda_kernel; + +use super::ops::{try_unary_op, UnaryKernel}; +use crate::{shapes::*, tensor::*}; + +#[repr(C)] +#[derive(Debug, Default, Copy, Clone)] +pub struct SigmoidrKernelOp; + +/// [Sigmoid](https://en.wikipedia.org/wiki/Sigmoid_function). `1 / (1 + exp(-t))`. +/// Basically the same as sigmoid but will always return non-zero gradients. +/// The derivative is `sigmoid(t) * (1.0 - sigmoid(t))`. +/// +/// Examples: +/// ```rust +/// # use dfdx::prelude::*; +/// # let dev: Cpu = Default::default(); +/// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]); +/// let r = t.sigmoid(); +/// ``` +pub fn sigmoidr, T: Tape>( + t: Tensor, +) -> Tensor { + t.sigmoidr() +} + +impl, T: Tape> Tensor { + /// See [sigmoidr] + pub fn sigmoidr(self) -> Self { + self.try_sigmoidr().unwrap() + } + /// See [sigmoidr] + pub fn try_sigmoidr(self) -> Result { + try_unary_op(SigmoidrKernelOp, self) + } +} + +#[cfg(test)] +mod tests { + use crate::{tensor::*, tensor_ops::*, tests::*}; + + #[test] + fn test_sigmoidr() { + let dev: TestDevice = Default::default(); + let x = dev + .tensor([-2.0, -1.0, 0.0, 1.0, -TestDtype::INFINITY]) + .to_dtype::(); + let r = x.leaky_trace().sigmoidr(); + assert_close_to_literal!(r, [0.11920292, 0.26894143, 0.5, 0.7310586, 0.0]); + let g = r.mean().backward(); + assert_close_to_literal!( + g.get(&x), + [0.020998716, 0.039322387, 0.05, 0.039322387, 0.00000002] + ); + } +} diff --git a/src/tensor_ops/sigmoidr/sigmoidr.cu b/src/tensor_ops/sigmoidr/sigmoidr.cu new file mode 100644 index 00000000..abc8798d --- /dev/null +++ b/src/tensor_ops/sigmoidr/sigmoidr.cu @@ -0,0 +1,29 @@ +#include "unary_op_macros.cuh" + +struct SigmoidrKernelOp {}; + +template +__device__ __forceinline__ T sigmoidr_fwd(T x) { + T one = 1.0; + return one / (one + expg(-x)); +} + +template +__device__ __forceinline__ T sigmoidr_bwd(T y) { + T one = 1.0; + T d = y * (one - y); + return max(d, 0.0000001); +} + +UNARY_OP(__half, sigmoidr_fwd_f16, sigmoidr_bwd_f16, SigmoidrKernelOp, + sigmoidr_fwd(x), + sigmoidr_bwd(y)) + +UNARY_OP(float, sigmoidr_fwd_f32, sigmoidr_bwd_f32, SigmoidrKernelOp, + sigmoidr_fwd(x), + sigmoidr_bwd(y)) + +UNARY_OP(double, sigmoidr_fwd_f64, sigmoidr_bwd_f64, SigmoidrKernelOp, + sigmoidr_fwd(x), + sigmoidr_bwd(y)) + \ No newline at end of file diff --git a/src/tensor_ops/utilities/device.rs b/src/tensor_ops/utilities/device.rs index 8a195d6a..55312e82 100644 --- a/src/tensor_ops/utilities/device.rs +++ b/src/tensor_ops/utilities/device.rs @@ -89,6 +89,7 @@ pub trait Device: + UnaryKernel + UnaryKernel + UnaryKernel + + UnaryKernel + UnaryKernel + UnaryKernel + UnaryKernel