-
Notifications
You must be signed in to change notification settings - Fork 0
/
kautodiff.h
246 lines (203 loc) · 11.2 KB
/
kautodiff.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
/*
The MIT License
Copyright (c) 2018-2019 Dana-Farber Cancer Institute
2016-2018 Broad Institute
Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:
The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/
#ifndef KANN_AUTODIFF_H
#define KANN_AUTODIFF_H
#define KAD_VERSION "r544"
#include <stdio.h>
#include <stdint.h>
#ifdef __STRICT_ANSI__
#define inline
#endif
#define KAD_MAX_DIM 4 /* max dimension */
#define KAD_MAX_OP 64 /* max number of operators */
/* A computational graph is a directed acyclic graph. In the graph, an external
* node represents a variable, a constant or a feed; an internal node
* represents an operator; an edge from node v to w indicates v is an operand
* of w.
*/
#define KAD_VAR 0x1
#define KAD_CONST 0x2
#define KAD_POOL 0x4
#define KAD_SHARE_RNG 0x10 /* with this flag on, different time step shares the same RNG status after unroll */
#define kad_is_back(p) ((p)->flag & KAD_VAR)
#define kad_is_ext(p) ((p)->n_child == 0)
#define kad_is_var(p) (kad_is_ext(p) && kad_is_back(p))
#define kad_is_const(p) (kad_is_ext(p) && ((p)->flag & KAD_CONST))
#define kad_is_feed(p) (kad_is_ext(p) && !kad_is_back(p) && !((p)->flag & KAD_CONST))
#define kad_is_pivot(p) ((p)->n_child == 1 && ((p)->flag & KAD_POOL))
#define kad_is_switch(p) ((p)->op == 12 && !((p)->flag & KAD_POOL))
#define kad_use_rng(p) ((p)->op == 15 || (p)->op == 24)
#define kad_eval_enable(p) ((p)->tmp = 1)
#define kad_eval_disable(p) ((p)->tmp = -1)
/* a node in the computational graph */
typedef struct kad_node_t {
uint8_t n_d; /* number of dimensions; no larger than KAD_MAX_DIM */
uint8_t flag; /* type of the node; see KAD_F_* for valid flags */
uint16_t op; /* operator; kad_op_list[op] is the actual function */
int32_t n_child; /* number of operands/child nodes */
int32_t tmp; /* temporary field; MUST BE zero before calling kad_compile() */
int32_t ptr_size; /* size of ptr below */
int32_t d[KAD_MAX_DIM]; /* dimensions */
int32_t ext_label; /* labels for external uses (not modified by the kad_* APIs) */
uint32_t ext_flag; /* flags for external uses (not modified by the kad_* APIs) */
float *x; /* value; allocated for internal nodes */
float *g; /* gradient; allocated for internal nodes */
void *ptr; /* for special operators that need additional parameters (e.g. conv2d) */
void *gtmp; /* temporary data generated at the forward pass but used at the backward pass */
struct kad_node_t **child; /* operands/child nodes */
struct kad_node_t *pre; /* usually NULL; only used for RNN */
} kad_node_t, *kad_node_p;
#ifdef __cplusplus
extern "C" {
#endif
/**
* Compile/linearize a computational graph
*
* @param n_node number of nodes (out)
* @param n_roots number of nodes without predecessors
* @param roots list of nodes without predecessors
*
* @return list of nodes, of size *n_node
*/
kad_node_t **kad_compile_array(int *n_node, int n_roots, kad_node_t **roots);
kad_node_t **kad_compile(int *n_node, int n_roots, ...); /* an alternative API to above */
void kad_delete(int n, kad_node_t **a); /* deallocate a compiled/linearized graph */
/**
* Compute the value at a node
*
* @param n number of nodes
* @param a list of nodes
* @param from compute the value at this node, 0<=from<n
*
* @return a pointer to the value (pointing to kad_node_t::x, so don't call
* free() on it!)
*/
const float *kad_eval_at(int n, kad_node_t **a, int from);
void kad_eval_marked(int n, kad_node_t **a);
int kad_sync_dim(int n, kad_node_t **v, int batch_size);
/**
* Compute gradient
*
* @param n number of nodes
* @param a list of nodes
* @param from the function node; must be a scalar (compute \nabla a[from])
*/
void kad_grad(int n, kad_node_t **a, int from);
/**
* Unroll a recurrent computation graph
*
* @param n_v number of nodes
* @param v list of nodes
* @param new_n number of nodes in the unrolled graph (out)
* @param len how many times to unroll, one for each pivot
*
* @return list of nodes in the unrolled graph
*/
kad_node_t **kad_unroll(int n_v, kad_node_t **v, int *new_n, int *len);
int kad_n_pivots(int n_v, kad_node_t **v);
kad_node_t **kad_clone(int n, kad_node_t **v, int batch_size);
/* define a variable, a constant or a feed (placeholder in TensorFlow) */
kad_node_t *kad_var(float *x, float *g, int n_d, ...); /* a variable; gradients to be computed; not unrolled */
kad_node_t *kad_const(float *x, int n_d, ...); /* a constant; no gradients computed; not unrolled */
kad_node_t *kad_feed(int n_d, ...); /* an input/output; no gradients computed; unrolled */
/* operators taking two operands */
kad_node_t *kad_add(kad_node_t *x, kad_node_t *y); /* f(x,y) = x + y (generalized element-wise addition; f[i*n+j]=x[i*n+j]+y[j], n=kad_len(y), 0<j<n, 0<i<kad_len(x)/n) */
kad_node_t *kad_sub(kad_node_t *x, kad_node_t *y); /* f(x,y) = x - y (generalized element-wise subtraction) */
kad_node_t *kad_mul(kad_node_t *x, kad_node_t *y); /* f(x,y) = x * y (generalized element-wise product) */
kad_node_t *kad_matmul(kad_node_t *x, kad_node_t *y); /* f(x,y) = x * y (general matrix product) */
kad_node_t *kad_cmul(kad_node_t *x, kad_node_t *y); /* f(x,y) = x * y^T (column-wise matrix product; i.e. y is transposed) */
/* loss functions; output scalar */
kad_node_t *kad_mse(kad_node_t *x, kad_node_t *y); /* mean square error */
kad_node_t *kad_ce_multi(kad_node_t *x, kad_node_t *y); /* multi-class cross-entropy; x is the preidction and y is the truth */
kad_node_t *kad_ce_bin(kad_node_t *x, kad_node_t *y); /* binary cross-entropy for (0,1) */
kad_node_t *kad_ce_bin_neg(kad_node_t *x, kad_node_t *y); /* binary cross-entropy for (-1,1) */
kad_node_t *kad_ce_multi_weighted(kad_node_t *pred, kad_node_t *truth, kad_node_t *weight);
#define KAD_PAD_NONE 0 /* use the smallest zero-padding */
#define KAD_PAD_SAME (-2) /* output to have the same dimension as input */
kad_node_t *kad_conv2d(kad_node_t *x, kad_node_t *w, int r_stride, int c_stride, int r_pad, int c_pad); /* 2D convolution with weight matrix flipped */
kad_node_t *kad_max2d(kad_node_t *x, int kernel_h, int kernel_w, int r_stride, int c_stride, int r_pad, int c_pad); /* 2D max pooling */
kad_node_t *kad_conv1d(kad_node_t *x, kad_node_t *w, int stride, int pad); /* 1D convolution with weight flipped */
kad_node_t *kad_max1d(kad_node_t *x, int kernel_size, int stride, int pad); /* 1D max pooling */
kad_node_t *kad_avg1d(kad_node_t *x, int kernel_size, int stride, int pad); /* 1D average pooling */
kad_node_t *kad_dropout(kad_node_t *x, kad_node_t *r); /* dropout at rate r */
kad_node_t *kad_sample_normal(kad_node_t *x); /* f(x) = x * r, where r is drawn from a standard normal distribution */
/* operators taking one operand */
kad_node_t *kad_square(kad_node_t *x); /* f(x) = x^2 (element-wise square) */
kad_node_t *kad_sigm(kad_node_t *x); /* f(x) = 1/(1+exp(-x)) (element-wise sigmoid) */
kad_node_t *kad_tanh(kad_node_t *x); /* f(x) = (1-exp(-2x)) / (1+exp(-2x)) (element-wise tanh) */
kad_node_t *kad_relu(kad_node_t *x); /* f(x) = max{0,x} (element-wise rectifier, aka ReLU) */
kad_node_t *kad_softmax(kad_node_t *x);/* f_i(x_1,...,x_n) = exp(x_i) / \sum_j exp(x_j) (softmax: tf.nn.softmax(x,dim=-1)) */
kad_node_t *kad_1minus(kad_node_t *x); /* f(x) = 1 - x */
kad_node_t *kad_exp(kad_node_t *x); /* f(x) = exp(x) */
kad_node_t *kad_log(kad_node_t *x); /* f(x) = log(x) */
kad_node_t *kad_sin(kad_node_t *x); /* f(x) = sin(x) */
kad_node_t *kad_stdnorm(kad_node_t *x); /* layer normalization; applied to the last dimension */
/* operators taking an indefinite number of operands (e.g. pooling) */
kad_node_t *kad_avg(int n, kad_node_t **x); /* f(x_1,...,x_n) = \sum_i x_i/n (mean pooling) */
kad_node_t *kad_max(int n, kad_node_t **x); /* f(x_1,...,x_n) = max{x_1,...,x_n} (max pooling) */
kad_node_t *kad_stack(int n, kad_node_t **x); /* f(x_1,...,x_n) = [x_1,...,x_n] (stack pooling) */
kad_node_t *kad_select(int n, kad_node_t **x, int which); /* f(x_1,...,x_n;i) = x_i (select pooling; -1 for the last) */
/* dimension reduction */
kad_node_t *kad_reduce_sum(kad_node_t *x, int axis); /* tf.reduce_sum(x, axis) */
kad_node_t *kad_reduce_mean(kad_node_t *x, int axis); /* tf.reduce_mean(x, axis) */
/* special operators */
kad_node_t *kad_slice(kad_node_t *x, int axis, int start, int end); /* take a slice on the axis-th dimension */
kad_node_t *kad_concat(int axis, int n, ...); /* concatenate on the axis-th dimension */
kad_node_t *kad_concat_array(int axis, int n, kad_node_t **p); /* the array version of concat */
kad_node_t *kad_reshape(kad_node_t *x, int n_d, int *d); /* reshape; similar behavior to TensorFlow's reshape() */
kad_node_t *kad_reverse(kad_node_t *x, int axis);
kad_node_t *kad_switch(int n, kad_node_t **p); /* manually (as a hyperparameter) choose one input, default to 0 */
/* miscellaneous operations on a compiled graph */
int kad_size_var(int n, kad_node_t *const* v); /* total size of all variables */
int kad_size_const(int n, kad_node_t *const* v); /* total size of all constants */
/* graph I/O */
int kad_save(FILE *fp, int n_node, kad_node_t **node);
kad_node_t **kad_load(FILE *fp, int *_n_node);
/* random number generator */
void *kad_rng(void);
void kad_srand(void *d, uint64_t seed);
uint64_t kad_rand(void *d);
double kad_drand(void *d);
double kad_drand_normal(void *d);
void kad_saxpy(int n, float a, const float *x, float *y);
/* debugging routines */
void kad_trap_fe(void); /* abort on divide-by-zero and NaN */
void kad_print_graph(FILE *fp, int n, kad_node_t **v);
void kad_check_grad(int n, kad_node_t **a, int from);
#ifdef __cplusplus
}
#endif
#define KAD_ALLOC 1
#define KAD_FORWARD 2
#define KAD_BACKWARD 3
#define KAD_SYNC_DIM 4
typedef int (*kad_op_f)(kad_node_t*, int);
extern kad_op_f kad_op_list[KAD_MAX_OP];
extern char *kad_op_name[KAD_MAX_OP];
static inline int kad_len(const kad_node_t *p) /* calculate the size of p->x */
{
int n = 1, i;
for (i = 0; i < p->n_d; ++i) n *= p->d[i];
return n;
}
#endif