custos_math/ops/nn/
activations.rs

1use crate::{each_op, each_op_slice_mut, Matrix};
2use custos::{impl_stack, number::Float, CDatatype, Device, MainMemory, Shape};
3
4#[cfg(feature = "cpu")]
5use custos::CPU;
6
7#[cfg(feature = "stack")]
8use custos::Stack;
9
10#[cfg(feature = "opencl")]
11use crate::{cl_str_op_mut, opencl::cl_str_op_mat};
12#[cfg(feature = "opencl")]
13use custos::OpenCL;
14
15#[cfg(feature = "cuda")]
16use crate::{cu_str_op, cu_str_op_mut};
17#[cfg(feature = "cuda")]
18use custos::CUDA;
19
20impl<'a, T, D: ActivationOps<T, S>, S: Shape> Matrix<'a, T, D, S> {
21    #[inline]
22    pub fn tanh(&self) -> Matrix<'a, T, D, S> {
23        self.device().tanh(self)
24    }
25
26    #[inline]
27    pub fn tanh_grad(&self) -> Matrix<'a, T, D, S> {
28        self.device().tanh_grad(self)
29    }
30
31    #[inline]
32    pub fn relu(&self) -> Matrix<'a, T, D, S> {
33        self.device().relu(self)
34    }
35
36    #[inline]
37    pub fn relu_mut(&mut self) {
38        self.device().relu_mut(self)
39    }
40
41    #[inline]
42    pub fn relu_grad(&self) -> Matrix<'a, T, D, S> {
43        self.device().relu_grad(self)
44    }
45
46    #[inline]
47    pub fn relu_grad_mut(&mut self) {
48        self.device().relu_grad_mut(self)
49    }
50
51    #[inline]
52    pub fn sigmoid(&self) -> Matrix<'a, T, D, S> {
53        self.device().sigmoid(self)
54    }
55
56    #[inline]
57    /// uses pre-computed sigmoid activation
58    pub fn sigmoid_grad(&self) -> Matrix<'a, T, D, S> {
59        self.device().sigmoid_grad(self)
60    }
61}
62
63pub trait ActivationOps<T, S: Shape = (), D: Device = Self>: Device {
64    fn sigmoid(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
65    fn sigmoid_grad(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
66    fn tanh(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
67    fn tanh_grad(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
68    fn relu(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
69    /// inplace
70    fn relu_mut(&self, x: &mut Matrix<T, D, S>);
71    fn relu_grad(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
72    /// inplace
73    fn relu_grad_mut(&self, x: &mut Matrix<T, D, S>);
74}
75
76#[cfg(feature = "opencl")]
77impl<T: CDatatype + Float> ActivationOps<T> for OpenCL {
78    #[inline]
79    fn sigmoid(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
80        cl_str_op_mat(self, x, "1.0 / (1.0 + exp(-x))").unwrap()
81    }
82
83    fn sigmoid_grad(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
84        cl_str_op_mat(self, x, "x * (1.0 - x)").unwrap()
85    }
86
87    #[inline]
88    fn tanh(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
89        cl_str_op_mat(self, x, "tanh(x)").unwrap()
90    }
91
92    #[inline]
93    fn tanh_grad(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
94        cl_str_op_mat(self, x, "1.0 - pow(tanh(x), 2)").unwrap()
95    }
96
97    #[inline]
98    fn relu(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
99        cl_str_op_mat(self, x, "x * (x >= 0)").unwrap()
100    }
101
102    #[inline]
103    fn relu_mut(&self, x: &mut Matrix<T, Self, ()>) {
104        cl_str_op_mut(self, x, "x * (x >= 0)").unwrap();
105    }
106
107    #[inline]
108    fn relu_grad(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
109        cl_str_op_mat(self, x, "(x >= 0)").unwrap()
110    }
111
112    #[inline]
113    fn relu_grad_mut(&self, x: &mut Matrix<T, Self, ()>) {
114        cl_str_op_mut(self, x, "(x >= 0)").unwrap()
115    }
116}
117
118#[impl_stack]
119impl<T: Float, D: MainMemory, S: Shape> ActivationOps<T, S, D> for CPU {
120    #[inline]
121    fn sigmoid(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
122        each_op(self, x, |x| T::one() / (T::one() + -x.exp()))
123    }
124    #[inline]
125    fn sigmoid_grad(&self, activated: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
126        each_op(self, activated, |x| x * (T::one() - x))
127    }
128
129    #[inline]
130    fn tanh(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
131        each_op(self, x, |x| x.tanh())
132    }
133
134    #[inline]
135    fn tanh_grad(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
136        each_op(self, x, |x| T::one() - x.tanh().powi(2))
137    }
138
139    #[inline]
140    fn relu(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
141        each_op(self, x, |x| T::from_usize((x >= T::zero()) as usize) * x)
142    }
143
144    #[inline]
145    fn relu_mut(&self, x: &mut Matrix<T, D, S>) {
146        each_op_slice_mut(x, |x| T::from_usize((x >= T::zero()) as usize) * x)
147    }
148
149    #[inline]
150    fn relu_grad(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
151        each_op(self, x, |x| T::from_usize((x >= T::default()) as usize))
152    }
153
154    #[inline]
155    fn relu_grad_mut(&self, x: &mut Matrix<T, D, S>) {
156        each_op_slice_mut(x, |x| T::from_usize((x >= T::default()) as usize))
157    }
158}
159
160#[cfg(feature = "cuda")]
161impl<T: CDatatype> ActivationOps<T> for CUDA {
162    #[inline]
163    fn sigmoid(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
164        let out = cu_str_op(self, x, "1.0 / (1.0 + exp(-x))").unwrap();
165        (out, x.dims()).into()
166    }
167
168    fn sigmoid_grad(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
169        let out = cu_str_op(self, x, "x * (1.0 - x)").unwrap();
170        (out, x.dims()).into()
171    }
172
173    #[inline]
174    fn tanh(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
175        let out = cu_str_op(self, x, "tanh(x)").unwrap();
176        (out, x.dims()).into()
177    }
178
179    #[inline]
180    fn tanh_grad(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
181        let out = cu_str_op(self, x, "1.0 - pow(tanh(x), 2)").unwrap();
182        (out, x.dims()).into()
183    }
184
185    #[inline]
186    fn relu(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
187        let out = cu_str_op(self, x, "x * (x >= 0)").unwrap();
188        (out, x.dims()).into()
189    }
190
191    #[inline]
192    fn relu_grad(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
193        let out = cu_str_op(self, x, "(x >= 0)").unwrap();
194        (out, x.dims()).into()
195    }
196
197    #[inline]
198    fn relu_mut(&self, x: &mut Matrix<T, Self, ()>) {
199        cu_str_op_mut(self, x, "x * (x >= 0)").unwrap();
200    }
201
202    #[inline]
203    fn relu_grad_mut(&self, x: &mut Matrix<T, Self, ()>) {
204        cu_str_op_mut(self, x, "(x >= 0)").unwrap()
205    }
206}
custos_math/ops/nn/activations.rs

custos_math/ops/nn/
activations.rs
