-
Notifications
You must be signed in to change notification settings - Fork 0
/
ops.py
169 lines (143 loc) · 7.05 KB
/
ops.py
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
import tensorflow as tf
# INPUT: x is K x M
# is_test is a boolean which decided whether copies are actually added
# N is the number of mc samples to add
# OUTPUT: y is K x N x M, with N copies along of x along that dimension
def add_mc_samples(x, N, name='AddMCSamples'):
with tf.variable_scope(name):
x = tf.expand_dims(x, axis=1)
y = tf.tile(x, [1, N, 1])
return y
# INPUT: x is K x N x M
# Performs dropout on the input, with the same M x W dropout mask across the N
# dimension. Then rehsapes to (MN x W) and computes x*W + b
# OUTPUT: y is K x N x output_size
# reg_loss is a scalar corresponding to a regularization loss
def dense_with_dropout(x, output_size, prob, weight_reg=0,
nonlinearity=(lambda x: x), is_test=False, stddev=0.05,
bias_start=0.0, name='DenseDropout'):
with tf.variable_scope(name):
# (K,N,M) = x.get_shape().as_list()
# test_noise_shape = [1,N,M]
x_shape = tf.shape(x)
K = x_shape[0]
N = x_shape[1]
M = x_shape[2]
M_int = x.get_shape().as_list()[2]
W = tf.get_variable("W", [M_int, output_size], tf.float32,
initializer=tf.random_normal_initializer(stddev=stddev))
b = tf.get_variable("b", [output_size],
initializer=tf.random_normal_initializer(mean=bias_start, stddev=stddev))
h1 = tf.nn.dropout(x, keep_prob=1-prob, noise_shape=[1, N, M])
h1 = tf.reshape(h1, (K*N, M))
y = nonlinearity(tf.matmul(h1, W) + b)
y = tf.reshape(y, (K, N, output_size))
reg_loss = weight_reg * tf.reduce_sum(W**2) / (1. - prob)
return y, reg_loss
def concrete_dropout(x, p, noise_shape=None, temp=0.1):
if not noise_shape:
noise_shape = tf.shape(x)
eps = 1e-9
unif_noise = tf.random_uniform(shape=noise_shape)
drop_prob = (
tf.log(p + eps)
- tf.log(1. - p + eps)
+ tf.log(unif_noise + eps)
- tf.log(1. - unif_noise + eps)
)
drop_prob = tf.sigmoid(drop_prob / temp)
drop_mask = 1. - drop_prob
retain_prob = 1. - p
x *= drop_mask
x /= retain_prob
return x
# Performs dropout on the input, with the same M x W dropout mask across the N
# dimension. Then rehsapes to (MN x W) and computes x*W + b
# OUTPUT: y is K x N x output_size
# reg_loss is a scalar corresponding to a regularization loss term
def dense_with_concrete_dropout(x, output_size, nonlinearity=(lambda x: x),
weight_reg = 0,
dropout_reg = 1e-5,
is_test=False, stddev=0.05, bias_start=0.0,
name='DenseConcreteDropout'):
with tf.variable_scope(name):
x_shape = tf.shape(x)
K = x_shape[0]
N = x_shape[1]
M = x_shape[2]
M_int = x.get_shape().as_list()[2]
W = tf.get_variable("W", [M_int, output_size], tf.float32,
initializer=tf.random_normal_initializer(stddev=stddev))
b = tf.get_variable("b", [output_size],
initializer=tf.constant_initializer(bias_start))
p_logit = tf.get_variable("p_logit", [], tf.float32,
initializer=tf.random_uniform_initializer(-2., 0.))
p = tf.sigmoid(p_logit)
tf.summary.scalar("p", p)
h1 = concrete_dropout(x, p, noise_shape=[1, N, M])
h1 = tf.reshape(h1, (-1, M))
y = nonlinearity(tf.matmul(h1, W) + b)
y = tf.reshape(y, (-1, N, output_size))
lengthscale = 1
kernel_regularizer = dropout_reg * lengthscale * tf.reduce_sum(W**2) / (1. - p) / 2
dropout_regularizer = M_int * dropout_reg * ( p * tf.log(p) + (1. - p) * tf.log(1. - p) )
reg_loss = kernel_regularizer + dropout_regularizer
return y, reg_loss
# Applies a sequence of dense_with_dropout layers
# INPUT: x, the input to the layer
# hidden_layer_sizes, the widths of the hidden layers
# dropout_prob, the probabality of dropping the input neurons
# is_test, boolean corresponding to the phase of operation
# weight_regularizer, the coefficient for regularizing the layer weights
# OUTPUT: current_input, the output of the MLP
# reg_loss, the regularization loss corresponding to all the layers
def mlp_with_dropout(x, hidden_layer_sizes, dropout_prob, is_test, weight_reg=0):
current_input = x
reg_losses = []
for i,s in enumerate(hidden_layer_sizes):
dropout = 0.0 if (i==0) else dropout_prob
current_input, reg_loss = dense_with_dropout(current_input, output_size=s,
prob=dropout, weight_reg=weight_reg,
nonlinearity=tf.nn.relu, is_test=is_test, name="fc%d"%(i))
reg_losses.append(reg_loss)
print("layer {}, shape {}".format(i, current_input.get_shape().as_list()))
reg_loss = tf.add_n(reg_losses)
return current_input, reg_loss
# Applies a sequence of dense_with_concrete_dropout layers
# INPUT: x, the input to the layer
# hidden_layer_sizes, the widths of the hidden layers
# is_test, boolean corresponding to the phase of operation
# weight_regularizer, the coefficient for regularizing the layer weights
# dropout_regularizer, the coefficient for regularizing the dropout probability
# OUTPUT: current_input, the output of the MLP
# reg_loss, the regularization loss corresponding to all the layers
def mlp_with_concrete_dropout(x, hidden_layer_sizes, is_test, weight_reg=0, dropout_reg=1e-5):
current_input = x
reg_losses = []
for i,s in enumerate(hidden_layer_sizes):
# if i == 0:
# current_input, reg_loss = dense_with_dropout(current_input, output_size=s,
# prob=0.0, weight_regularizer=weight_regularizer,
# nonlinearity=tf.nn.relu, is_test=is_test, name="fc%d"%(i))
# else:
current_input, reg_loss = dense_with_concrete_dropout(current_input, output_size=s,
weight_reg=weight_reg,
dropout_reg=dropout_reg,
nonlinearity=tf.nn.relu, is_test=is_test, name="fc%d"%(i))
reg_losses.append(reg_loss)
print("layer {}, shape {}".format(i, current_input.get_shape().as_list()))
reg_loss = tf.add_n(reg_losses)
return current_input, reg_loss
# INPUT: all inputs are K x N x M
# OUTPUT: K x N array of distances between x_hat and x in the mahalanobis sense
def mahalanobis_distance(x, x_hat, var, name="MahalanobisDist"):
eps = 1e-7
with tf.variable_scope(name):
return tf.reduce_sum( (x - x_hat)**2 / (var+eps), axis=2 )
# INPUT: x has size (K x N x M)
# OUTPUT: (K x M) array of the variance along the 2nd dimension
def epistemic_unc(x, name="EpistemicUnc"):
with tf.variable_scope(name):
mean = tf.reduce_mean( x, axis=1, keep_dims=True )
var = tf.reduce_mean( (x - mean)**2, axis=1 )
return var