Coverage for mlair/model_modules/probability_models.py: 0%

1""" 

2>>> MyCustomisedModel().model.compile(**kwargs) == MyCustomisedModel().compile(**kwargs) 

3True 

4 

5""" 

6 

7import mlair.model_modules.keras_extensions 

8 

9__author__ = "Felix Kleinert" 

10__date__ = '2022-07-08' 

11 

12import tensorflow as tf 

13import tensorflow.keras as keras 

14import tensorflow_probability as tfp 

15tfd = tfp.distributions 

16tfb = tfp.bijectors 

17tfpl = tfp.layers 

18 

19import logging 

20from mlair.model_modules import AbstractModelClass 

21from mlair.model_modules.inception_model import InceptionModelBase 

22from mlair.model_modules.flatten import flatten_tail 

23from mlair.model_modules.advanced_paddings import PadUtils, Padding2D, SymmetricPadding2D 

24from mlair.model_modules.loss import l_p_loss 

25 

26 

27class MyUnetProb(AbstractModelClass): 

28 def __init__(self, input_shape: list, output_shape: list, num_of_training_samples: int): 

29 super().__init__(input_shape[0], output_shape[0]) 

30 self.first_filter_size = 16 # 16*2#self._input_shape[-1] # 16 

31 self.lstm_units = 64 * 2 # * 2 

32 self.kernel_size = (3, 1) # (3,1) 

33 self.activation = "elu" 

34 self.pool_size = (2, 1) 

35 

36 self.num_of_training_samples = num_of_training_samples 

37 # self.divergence_fn = lambda q, p, _: tfd.kl_divergence(q, p) / input_shape[0][0] 

38 self.divergence_fn = lambda q, p, _: tfd.kl_divergence(q, p) / num_of_training_samples 

39 

40 # self.loss_fn = lambda y_true, y_pred: -y_pred.log_prob(y_true) 

41 # self.loss = nll 

42 

43 self.dropout = .15 # .2 

44 self.k_mixed_components = 2 

45 self.kernel_regularizer = keras.regularizers.l1_l2(l1=0.01, l2=0.01) 

46 self.bias_regularizer = keras.regularizers.l1_l2(l1=0.01, l2=0.01) 

47 

48 self.kernel_initializer = 'he_normal' 

49 

50 self.dense_units = 32 * 2 

51 self.initial_lr = 0.001 

52 

53 # apply to model 

54 self.set_model() 

55 self.set_compile_options() 

56 self.set_custom_objects(SymmetricPadding2D=SymmetricPadding2D, loss=self.loss, divergence_fn=self.divergence_fn) 

57 

58 def set_model(self): 

59 input_train = keras.layers.Input(shape=self._input_shape) 

60 pad_size = PadUtils.get_padding_for_same(self.kernel_size) 

61 

62 c1 = Padding2D("SymPad2D")(padding=pad_size)(input_train) 

63 c1 = tfpl.Convolution2DReparameterization( 

64 self.first_filter_size, self.kernel_size, padding='valid', 

65 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

66 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

67 kernel_divergence_fn=self.divergence_fn, 

68 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

69 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

70 bias_divergence_fn=self.divergence_fn, 

71 activation=self.activation, 

72 )(c1) 

73 #c1 = keras.layers.Conv2D(self.first_filter_size, self.kernel_size, activation=self.activation, 

74 # kernel_initializer=self.kernel_initializer, kernel_regularizer=self.kernel_regularizer, 

75 # bias_regularizer=self.bias_regularizer)(c1) 

76 c1 = keras.layers.Dropout(self.dropout)(c1) 

77 c1 = Padding2D("SymPad2D")(padding=pad_size)(c1) 

78 c1 = tfpl.Convolution2DReparameterization( 

79 self.first_filter_size, self.kernel_size, padding='valid', 

80 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

81 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

82 kernel_divergence_fn=self.divergence_fn, 

83 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

84 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

85 bias_divergence_fn=self.divergence_fn, 

86 activation=self.activation, 

87 name='c1' 

88 )(c1) 

89 #c1 = keras.layers.Conv2D(self.first_filter_size, self.kernel_size, activation=self.activation, 

90 # kernel_initializer=self.kernel_initializer, name='c1', 

91 # kernel_regularizer=self.kernel_regularizer, 

92 # bias_regularizer=self.bias_regularizer)(c1) 

93 p1 = c1 

94 # p1 = keras.layers.MaxPooling2D(self.pool_size)(c1) 

95 

96 c2 = Padding2D("SymPad2D")(padding=pad_size)(p1) 

97 c2 = tfpl.Convolution2DReparameterization( 

98 self.first_filter_size * 2, self.kernel_size, padding='valid', 

99 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

100 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

101 kernel_divergence_fn=self.divergence_fn, 

102 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

103 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

104 bias_divergence_fn=self.divergence_fn, 

105 activation=self.activation, 

106 )(c2) 

107 # c2 = keras.layers.Conv2D(self.first_filter_size * 2, self.kernel_size, activation=self.activation, 

108 # kernel_initializer=self.kernel_initializer, kernel_regularizer=self.kernel_regularizer, 

109 # bias_regularizer=self.bias_regularizer)(c2) 

110 c2 = keras.layers.Dropout(self.dropout)(c2) 

111 c2 = Padding2D("SymPad2D")(padding=pad_size)(c2) 

112 c2 = tfpl.Convolution2DReparameterization( 

113 self.first_filter_size * 2, self.kernel_size, padding='valid', 

114 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

115 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

116 kernel_divergence_fn=self.divergence_fn, 

117 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

118 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

119 bias_divergence_fn=self.divergence_fn, 

120 activation=self.activation, name="c2" 

121 )(c2) 

122 # c2 = keras.layers.Conv2D(self.first_filter_size * 2, self.kernel_size, activation=self.activation, 

123 # kernel_initializer=self.kernel_initializer, name='c2', 

124 # kernel_regularizer=self.kernel_regularizer, 

125 # bias_regularizer=self.bias_regularizer)(c2) 

126 p2 = c2 

127 # p2 = keras.layers.MaxPooling2D(self.pool_size)(c2) 

128 

129 c3 = Padding2D("SymPad2D")(padding=pad_size)(p2) 

130 c3 = tfpl.Convolution2DReparameterization( 

131 self.first_filter_size * 4, self.kernel_size, padding='valid', 

132 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

133 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

134 kernel_divergence_fn=self.divergence_fn, 

135 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

136 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

137 bias_divergence_fn=self.divergence_fn, 

138 activation=self.activation 

139 )(c3) 

140 # c3 = keras.layers.Conv2D(self.first_filter_size * 4, self.kernel_size, activation=self.activation, 

141 # kernel_initializer=self.kernel_initializer, kernel_regularizer=self.kernel_regularizer, 

142 # bias_regularizer=self.bias_regularizer)(c3) 

143 c3 = keras.layers.Dropout(self.dropout * 2)(c3) 

144 c3 = Padding2D("SymPad2D")(padding=pad_size)(c3) 

145 c3 = tfpl.Convolution2DReparameterization( 

146 self.first_filter_size * 4, self.kernel_size, padding='valid', 

147 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

148 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

149 kernel_divergence_fn=self.divergence_fn, 

150 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

151 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

152 bias_divergence_fn=self.divergence_fn, 

153 activation=self.activation, name="c3" 

154 )(c3) 

155 # c3 = keras.layers.Conv2D(self.first_filter_size * 4, self.kernel_size, activation=self.activation, 

156 # kernel_initializer=self.kernel_initializer, name='c3', 

157 # kernel_regularizer=self.kernel_regularizer, 

158 # bias_regularizer=self.bias_regularizer)(c3) 

159 # p3 = c3 

160 p3 = keras.layers.MaxPooling2D(self.pool_size)(c3) 

161 

162 ### own LSTM Block ### 

163 ls1 = keras.layers.Reshape((p3.shape[1], p3.shape[-1]))(p3) 

164 ls1 = keras.layers.LSTM(self.lstm_units, return_sequences=True)(ls1) 

165 ls1 = keras.layers.LSTM(self.lstm_units, return_sequences=True)(ls1) 

166 c4 = keras.layers.Reshape((p3.shape[1], 1, -1))(ls1) 

167 

168 ### own 2nd LSTM Block ### 

169 ls2 = keras.layers.Reshape((c3.shape[1], c3.shape[-1]))(c3) 

170 ls2 = keras.layers.LSTM(self.lstm_units, return_sequences=True)(ls2) 

171 ls2 = keras.layers.LSTM(self.lstm_units, return_sequences=True)(ls2) 

172 c4_2 = keras.layers.Reshape((c3.shape[1], 1, -1))(ls2) 

173 

174 u7 = keras.layers.UpSampling2D(size=(3, 1))(c4) 

175 cn3 = Padding2D("SymPad2D")(padding=pad_size)(c3) 

176 # u7 = c4 

177 u7 = keras.layers.concatenate([u7, cn3], name="u7_c3") 

178 c7 = u7 

179 # c7 = Padding2D("SymPad2D")(padding=pad_size)(u7) 

180 c7 = tfpl.Convolution2DReparameterization( 

181 self.first_filter_size * 4, self.kernel_size, padding='valid', 

182 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

183 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

184 kernel_divergence_fn=self.divergence_fn, 

185 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

186 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

187 bias_divergence_fn=self.divergence_fn, 

188 activation=self.activation 

189 )(c7) 

190 # c7 = keras.layers.Conv2D(self.first_filter_size * 4, self.kernel_size, activation=self.activation, 

191 # kernel_initializer=self.kernel_initializer, kernel_regularizer=self.kernel_regularizer, 

192 # bias_regularizer=self.bias_regularizer)(c7) 

193 c7 = keras.layers.concatenate([c7, c4_2], name="Concat_2nd_LSTM") 

194 c7 = keras.layers.Dropout(self.dropout * 2)(c7) 

195 c7 = Padding2D("SymPad2D")(padding=pad_size)(c7) 

196 c7 = tfpl.Convolution2DReparameterization( 

197 self.first_filter_size * 4, self.kernel_size, padding='valid', 

198 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

199 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

200 kernel_divergence_fn=self.divergence_fn, 

201 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

202 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

203 bias_divergence_fn=self.divergence_fn, 

204 activation=self.activation, name='c7_to_u8' 

205 )(c7) 

206 # c7 = keras.layers.Conv2D(self.first_filter_size * 4, self.kernel_size, activation=self.activation, 

207 # kernel_initializer=self.kernel_initializer, name='c7_to_u8', 

208 # kernel_regularizer=self.kernel_regularizer, 

209 # bias_regularizer=self.bias_regularizer)(c7) 

210 

211 

212 # u8 = Padding2D("SymPad2D")(padding=pad_size)(c7) 

213 # u8 = keras.layers.Conv2DTranspose(32, self.pool_size, strides=self.pool_size)(u8) 

214 u8 = c7 

215 # u8 = c3 

216 u8 = keras.layers.concatenate([u8, c2], name="u8_c2") 

217 c8 = Padding2D("SymPad2D")(padding=pad_size)(u8) 

218 c8 = tfpl.Convolution2DReparameterization( 

219 self.first_filter_size * 2, self.kernel_size, padding='valid', 

220 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

221 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

222 kernel_divergence_fn=self.divergence_fn, 

223 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

224 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

225 bias_divergence_fn=self.divergence_fn, 

226 activation=self.activation 

227 )(c8) 

228 # c8 = keras.layers.Conv2D(self.first_filter_size * 2, self.kernel_size, activation=self.activation, 

229 # kernel_initializer=self.kernel_initializer, kernel_regularizer=self.kernel_regularizer, 

230 # bias_regularizer=self.bias_regularizer)(c8) 

231 c8 = keras.layers.Dropout(self.dropout)(c8) 

232 c8 = Padding2D("SymPad2D")(padding=pad_size)(c8) 

233 c8 = tfpl.Convolution2DReparameterization( 

234 self.first_filter_size * 2, self.kernel_size, padding='valid', 

235 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

236 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

237 kernel_divergence_fn=self.divergence_fn, 

238 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

239 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

240 bias_divergence_fn=self.divergence_fn, 

241 activation=self.activation, name='c8_to_u9' 

242 )(c8) 

243 # c8 = keras.layers.Conv2D(self.first_filter_size * 2, self.kernel_size, activation=self.activation, 

244 # kernel_initializer=self.kernel_initializer, name='c8_to_u9', 

245 # kernel_regularizer=self.kernel_regularizer, 

246 # bias_regularizer=self.bias_regularizer)(c8) 

247 

248 # u9 = Padding2D("SymPad2D")(padding=pad_size)(c8) 

249 # u9 = keras.layers.Conv2DTranspose(16, self.pool_size, strides=self.pool_size)(u9) 

250 u9 = c8 

251 u9 = keras.layers.concatenate([u9, c1], name="u9_c1") 

252 c9 = Padding2D("SymPad2D")(padding=pad_size)(u9) 

253 c9 = tfpl.Convolution2DReparameterization( 

254 self.first_filter_size, self.kernel_size, padding='valid', 

255 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

256 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

257 kernel_divergence_fn=self.divergence_fn, 

258 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

259 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

260 bias_divergence_fn=self.divergence_fn, 

261 activation=self.activation, 

262 )(c9) 

263 # c9 = keras.layers.Conv2D(self.first_filter_size, self.kernel_size, activation=self.activation, 

264 # kernel_initializer=self.kernel_initializer, kernel_regularizer=self.kernel_regularizer, 

265 # bias_regularizer=self.bias_regularizer)(c9) 

266 c9 = keras.layers.Dropout(self.dropout)(c9) 

267 c9 = Padding2D("SymPad2D")(padding=pad_size)(c9) 

268 c9 = tfpl.Convolution2DReparameterization( 

269 self.first_filter_size, self.kernel_size, padding='valid', 

270 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

271 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

272 kernel_divergence_fn=self.divergence_fn, 

273 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

274 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

275 bias_divergence_fn=self.divergence_fn, 

276 activation=self.activation, 

277 )(c9) 

278 # c9 = keras.layers.Conv2D(self.first_filter_size, self.kernel_size, activation=self.activation, 

279 # kernel_initializer=self.kernel_initializer, name='c9', 

280 # kernel_regularizer=self.kernel_regularizer, 

281 # bias_regularizer=self.bias_regularizer)(c9) 

282 

283 # outputs = keras.layers.Conv2D(1, (1, 1), activation='sigmoid')(c9) 

284 dl = keras.layers.Flatten()(c9) 

285 dl = keras.layers.Dropout(self.dropout)(dl) 

286 

287 # outputs = tfpl.DenseVariational(tfpl.MultivariateNormalTriL.params_size(self._output_shape), 

288 # make_posterior_fn=self.posterior, 

289 # make_prior_fn=self.prior)(dl) 

290 # outputs = tfpl.MultivariateNormalTriL(self._output_shape)(outputs) 

291 # outputs = keras.layers.Dense(units=self._output_shape)(dl) 

292 

293 #outputs = keras.layers.Dense(tfpl.IndependentNormal.params_size(self._output_shape), 

294 # )(dl) 

295 #outputs = tfpl.DenseVariational(units=tfpl.IndependentNormal.params_size(self._output_shape), 

296 # #make_prior_fn=self.prior, 

297 # make_prior_fn=prior_trainable, 

298 # make_posterior_fn=self.posterior, 

299 # )(dl) 

300 #outputs = VarDense(units=tfpl.IndependentNormal.params_size(self._output_shape), 

301 # make_prior_fn=self.prior, 

302 # make_posterior_fn=self.posterior, 

303 # )(dl) 

304 

305 

306 #outputs = tfpl.IndependentNormal(self._output_shape)(outputs) 

307 params_size = tfpl.MixtureSameFamily.params_size( 

308 self.k_mixed_components, 

309 component_params_size=tfpl.MultivariateNormalTriL.params_size(self._output_shape) 

310 ) 

311 

312 pars = tf.keras.layers.Dense(params_size)(dl) 

313 # pars = DenseVariationalCustom( 

314 # units=params_size, make_prior_fn=prior, make_posterior_fn=posterior, 

315 # kl_use_exact=True, kl_weight=1./self.x_train_shape)(dl) 

316 

317 outputs = tfpl.MixtureSameFamily(self.k_mixed_components, 

318 tfpl.MultivariateNormalTriL( 

319 self._output_shape, 

320 convert_to_tensor_fn=tfp.distributions.Distribution.mode 

321 ) 

322 )(pars) 

323 

324 self.model = keras.Model(inputs=input_train, outputs=outputs) 

325 

326 def set_compile_options(self): 

327 # self.optimizer = keras.optimizers.Adam(lr=self.initial_lr, 

328 # clipnorm=self.clipnorm, 

329 # ) 

330 

331 # loss = nll 

332 # self.optimizer = tf.keras.optimizers.RMSprop(learning_rate=self.initial_lr) 

333 self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.initial_lr) 

334 self.loss = nll 

335 self.compile_options = {"metrics": ["mse", "mae"]} 

336 

337 # loss = keras.losses.MeanSquaredError() 

338 # self.compile_options = {"loss": [loss]} 

339 

340 @staticmethod 

341 def prior(kernel_size, bias_size, dtype=None): 

342 n = kernel_size + bias_size 

343 

344 prior_model = tf.keras.Sequential([ 

345 

346 tfpl.DistributionLambda( 

347 # Note: Our prior is a non-trianable distribution 

348 lambda t: tfd.MultivariateNormalDiag(loc=tf.zeros(n), scale_diag=tf.ones(n))) 

349 ]) 

350 

351 return prior_model 

352 

353 @staticmethod 

354 def posterior(kernel_size, bias_size, dtype=None): 

355 n = kernel_size + bias_size 

356 

357 posterior_model = tf.keras.Sequential([ 

358 

359 tfpl.VariableLayer(tfpl.MultivariateNormalTriL.params_size(n), dtype=dtype), 

360 tfpl.MultivariateNormalTriL(n) 

361 ]) 

362 

363 return posterior_model 

364 

365 

366 

367 

368class MyCNNProb(AbstractModelClass): 

369 """ 

370 Taken fromhttps://towardsdatascience.com/uncertainty-in-deep-learning-bayesian-cnn-tensorflow-probability-758d7482bef6 

371 and modified to our data 

372 """ 

373 def __init__(self, input_shape: list, output_shape: list): 

374 super().__init__(input_shape[0], output_shape[0]) 

375 self.initial_lr = 0.001 

376 

377 self.divergence_fn = lambda q, p, q_tensor : self.approximate_kl(q, p, q_tensor) / 1000 # check how to get num of samples included here 

378 # apply to model 

379 self.set_model() 

380 self.set_compile_options() 

381 self.set_custom_objects(loss_fn=self.loss_fn ) 

382 

383 

384 @staticmethod 

385 def loss_fn(y_true, y_pred): 

386 return -y_pred.log_prob(y_true) 

387 

388 

389 # For Reparameterization Layers 

390 

391 @staticmethod 

392 def custom_normal_prior(dtype, shape, name, trainable, add_variable_fn): 

393 distribution = tfd.Normal(loc = 0.1 * tf.ones(shape, dtype), 

394 scale = 1.5 * tf.ones(shape, dtype)) 

395 batch_ndims = tf.size(distribution.batch_shape_tensor()) 

396 

397 distribution = tfd.Independent(distribution, 

398 reinterpreted_batch_ndims = batch_ndims) 

399 return distribution 

400 

401 @staticmethod 

402 def laplace_prior(dtype, shape, name, trainable, add_variable_fn): 

403 distribution = tfd.Laplace(loc = tf.zeros(shape, dtype), 

404 scale = tf.ones(shape, dtype)) 

405 batch_ndims = tf.size(distribution.batch_shape_tensor()) 

406 

407 distribution = tfd.Independent(distribution, 

408 reinterpreted_batch_ndims = batch_ndims) 

409 return distribution 

410 

411 

412 @staticmethod 

413 def approximate_kl(q, p, q_tensor): 

414 return tf.reduce_mean(q.log_prob(q_tensor) - p.log_prob(q_tensor)) 

415 

416 

417 def conv_reparameterization_layer(self, filters, kernel_size, activation): 

418 # For simplicity, we use default prior and posterior. 

419 # In the next parts, we will use custom mixture prior and posteriors. 

420 return tfpl.Convolution2DReparameterization( 

421 filters = filters, 

422 kernel_size = kernel_size, 

423 activation = activation, 

424 padding = 'same', 

425 kernel_posterior_fn = tfpl.default_mean_field_normal_fn(is_singular=False), 

426 kernel_prior_fn = tfpl.default_multivariate_normal_fn, 

427 

428 bias_prior_fn = tfpl.default_multivariate_normal_fn, 

429 bias_posterior_fn = tfpl.default_mean_field_normal_fn(is_singular=False), 

430 

431 kernel_divergence_fn = self.divergence_fn, 

432 bias_divergence_fn = self.divergence_fn) 

433 

434 def set_model(self): 

435 bayesian_cnn = tf.keras.Sequential([ 

436 tf.keras.layers.InputLayer(self._input_shape), 

437 self.conv_reparameterization_layer(16, 3, 'swish'), 

438 #tf.keras.layers.MaxPooling2D(2), 

439 self.conv_reparameterization_layer(32, 3, 'swish'), 

440 #tf.keras.layers.MaxPooling2D(2), 

441 self.conv_reparameterization_layer(64, 3, 'swish'), 

442 #tf.keras.layers.MaxPooling2D(2), 

443 self.conv_reparameterization_layer(128, 3, 'swish'), 

444 #tf.keras.layers.GlobalMaxPooling2D(), 

445 tfpl.DenseReparameterization( 

446 units=tfpl.IndependentNormal.params_size(self._output_shape), activation=None, 

447 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False), 

448 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

449 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

450 bias_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False), 

451 kernel_divergence_fn=self.divergence_fn, 

452 bias_divergence_fn=self.divergence_fn), 

453 tfpl.IndependentNormal(self._output_shape) 

454 ]) 

455 

456 input_train = keras.layers.Input(shape=self._input_shape) 

457 x = self.conv_reparameterization_layer(16, 3, 'swish')(input_train) 

458 # tf.keras.layers.MaxPooling2D(2), 

459 x = self.conv_reparameterization_layer(32, 3, 'swish')(x) 

460 # tf.keras.layers.MaxPooling2D(2), 

461 x = self.conv_reparameterization_layer(64, 3, 'swish')(x) 

462 # tf.keras.layers.MaxPooling2D(2), 

463 x = self.conv_reparameterization_layer(128, 3, 'swish')(x) 

464 x = tf.keras.layers.Flatten()(x) 

465 # x = tfpl.DenseReparameterization( 

466 # units=tfpl.IndependentNormal.params_size(self._output_shape), activation=None, 

467 # kernel_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False), 

468 # kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

469 # bias_prior_fn=tfpl.default_multivariate_normal_fn, 

470 # bias_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False), 

471 # kernel_divergence_fn=self.divergence_fn, 

472 # bias_divergence_fn=self.divergence_fn)(x) 

473 # outputs = tfpl.IndependentNormal(self._output_shape)(x) 

474 x = tf.keras.layers.Dense(tfpl.IndependentNormal.params_size(event_shape=self._output_shape))(x) 

475 outputs = tfpl.IndependentNormal(event_shape=self._output_shape)(x) 

476 # outputs = tfpl.DistributionLambda( 

477 # make_distribution_fn=lambda t: tfd.Normal( 

478 # loc=t[..., 0], scale=tf.exp(t[..., 1])), 

479 # convert_to_tensor_fn=lambda s: s.sample(30))(x) 

480 

481 

482 bnn = keras.Model(inputs=input_train, outputs=outputs) 

483 self.model = bnn 

484 

485 

486 logging.info(f"model summary:\n{self.model.summary()}") 

487 

488 def set_compile_options(self): 

489 self.optimizer = tf.keras.optimizers.Adam(lr=self.initial_lr, 

490 # clipnorm=self.clipnorm, 

491 ) 

492 

493 loss = self.loss_fn 

494 # self.compile_options = {"loss": [loss], "metrics": ["mse", "mae"]} 

495 

496 # loss = keras.losses.MeanSquaredError() 

497 self.compile_options = {"loss": [loss]} 

498 

499 

500 

501 

502class VarDense(tf.keras.layers.Layer): 

503 

504 def __init__(self, 

505 units, 

506 make_posterior_fn, 

507 make_prior_fn, 

508 kl_weight=None, 

509 kl_use_exact=False, 

510 activation=None, 

511 use_bias=True, 

512 activity_regularizer=None, 

513 **kwargs 

514 ): 

515 super().__init__(**kwargs) 

516 self.units = units 

517 self.make_posterior_fn = make_posterior_fn 

518 self.make_prior_fn = make_prior_fn 

519 self.kl_weight = kl_weight, 

520 self.kl_use_exact = kl_use_exact, 

521 self.activation = activation, 

522 self.use_bias = use_bias, 

523 self.activity_regularizer = activity_regularizer 

524 self.tfpllayer = tfpl.DenseVariational(units=self.units, 

525 make_prior_fn=self.make_prior_fn, 

526 make_posterior_fn=self.make_posterior_fn, 

527 kl_weight=self.kl_weight, 

528 kl_use_exact=self.kl_use_exact, 

529 use_bias=self.use_bias, 

530 activity_regularizer=self.activity_regularizer 

531 ) 

532 

533 def call(self, inputs): 

534 return self.tfpllayer(inputs) 

535 

536 

537 

538 

539 def get_config(self): 

540 config = super().get_config().copy() 

541 config.update({ 

542 "units": self.units, 

543 "make_posterior_fn": self.make_posterior_fn, 

544 "make_prior_fn": self.make_prior_fn, 

545 "kl_weight": self.kl_weight, 

546 "kl_use_exact": self.kl_use_exact, 

547 "activation": self.activation, 

548 "use_bias": self.use_bias, 

549 "activity_regularizer": self.activity_regularizer, 

550 }) 

551 return config 

552 

553 

554def prior_trainable(kernel_size, bias_size=0, dtype=None): 

555 n = kernel_size + bias_size 

556 return tf.keras.Sequential([ 

557 tfp.layers.VariableLayer(n, dtype=dtype), 

558 tfp.layers.DistributionLambda(lambda t: tfd.Independent( 

559 tfd.Normal(loc=t, scale=1), 

560 reinterpreted_batch_ndims=1)), 

561 ]) 

562 

563 

564class ProbTestModel(AbstractModelClass): 

565 def __init__(self, input_shape: list, output_shape: list): 

566 super().__init__(input_shape[0], output_shape[0]) 

567 self.initial_lr = 0.001 

568 self.loss = nll 

569 

570 self.set_model() 

571 self.set_compile_options() 

572 self.set_custom_objects(nll=nll) 

573 

574 def set_model(self): 

575 

576 x_in = keras.layers.Input(self._input_shape) 

577 x = keras.layers.Conv2D(kernel_size=(3,1), filters=8, 

578 activation='relu', padding="same")(x_in) 

579 x = keras.layers.Flatten()(x) 

580 x = keras.layers.Dense(tfpl.IndependentNormal.params_size(self._output_shape))(x) 

581 out = tfpl.IndependentNormal(self._output_shape)(x) 

582 model = keras.Model(inputs=x_in, outputs=out) 

583 

584 

585 #model = tf.keras.Sequential([ 

586 # keras.layers.InputLayer(self._input_shape), 

587 # keras.layers.Conv2D(kernel_size=(3,1), filters=8, 

588 # activation='relu', padding="same"), 

589 

590 # keras.layers.Flatten(), 

591 

592 # keras.layers.Dense(tfpl.IndependentNormal.params_size(self._output_shape)), 

593 # tfpl.IndependentNormal(self._output_shape, 

594 # convert_to_tensor_fn=tfp.distributions.Distribution.sample 

595 # ) 

596 

597 #]) 

598 self.model = model 

599 logging.info(self.model.summary()) 

600 

601 def set_compile_options(self): 

602 self.optimizer = tf.keras.optimizers.RMSprop(lr=self.initial_lr, 

603 # clipnorm=self.clipnorm, 

604 ) 

605 

606class ProbTestModel2(AbstractModelClass): 

607 def __init__(self, input_shape: list, output_shape: list): 

608 super().__init__(input_shape[0], output_shape[0]) 

609 self.initial_lr = 0.001 

610 self.loss = nll 

611 self.divergence_fn = lambda q, p, _: tfd.kl_divergence(q, p) / input_shape[0][0] 

612 

613 self.set_model() 

614 self.set_compile_options() 

615 self.set_custom_objects(nll=nll) 

616 

617 def set_model(self): 

618 model = tf.keras.Sequential([ 

619 tf.keras.layers.InputLayer(self._input_shape), 

620 #tf.keras.layers.Conv2D(kernel_size=(3,1), filters=8, 

621 # activation='relu', padding="same"), 

622 Convolution2DReparameterizationCustom( 

623 8, (3,1), padding='same', 

624 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

625 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

626 kernel_divergence_fn=self.divergence_fn, 

627 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

628 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

629 bias_divergence_fn=self.divergence_fn, 

630 activation='relu'), 

631 

632 tf.keras.layers.Flatten(), 

633 

634 tf.keras.layers.Dense(tfpl.MultivariateNormalTriL.params_size(self._output_shape)), 

635 tfpl.MultivariateNormalTriL(self._output_shape, 

636 convert_to_tensor_fn=tfp.distributions.Distribution.mode 

637 ) 

638 

639 ]) 

640 self.model = model 

641 logging.info(self.model.summary()) 

642 

643 def set_compile_options(self): 

644 self.optimizer = tf.keras.optimizers.RMSprop(lr=self.initial_lr, 

645 # clipnorm=self.clipnorm, 

646 ) 

647 

648 

649 

650 

651class ProbTestModel3(AbstractModelClass): 

652 def __init__(self, input_shape: list, output_shape: list): 

653 super().__init__(input_shape[0], output_shape[0]) 

654 

655 self.x_train_shape=100. 

656 self.set_model() 

657 self.set_compile_options() 

658 self.set_custom_objects(nll=nll) 

659 

660 def set_model(self): 

661 model = tf.keras.Sequential([ 

662 keras.layers.Flatten(input_shape=self._input_shape), 

663 # Epistemic uncertainty 

664 tfpl.DenseVariational(units=8, 

665 make_prior_fn=prior, 

666 make_posterior_fn=posterior, 

667 kl_weight=1/self.x_train_shape, 

668 kl_use_exact=False, 

669 activation='sigmoid'), 

670 

671 tfpl.DenseVariational(units=tfpl.IndependentNormal.params_size(1), 

672 make_prior_fn=prior, 

673 make_posterior_fn=posterior, 

674 kl_use_exact=False, 

675 kl_weight=1/self.x_train_shape), 

676 

677 # Aleatoric uncertainty 

678 tfpl.IndependentNormal(1) 

679 ]) 

680 logging.warning(model.summary()) 

681 self.model = model 

682 

683 def set_compile_options(self): 

684 self.optimizer = tf.keras.optimizers.RMSprop() 

685 

686 

687class ProbTestModel4(AbstractModelClass): 

688 def __init__(self, input_shape: list, output_shape: list): 

689 super().__init__(input_shape[0], output_shape[0]) 

690 

691 self.x_train_shape = 100. 

692 self.set_model() 

693 self.set_compile_options() 

694 self.set_custom_objects(nll=nll) 

695 

696 def set_model(self): 

697 model = tf.keras.Sequential([ 

698 keras.layers.Flatten(input_shape=self._input_shape), 

699 # Epistemic uncertainty 

700 DenseVariationalCustom(units=8, 

701 make_prior_fn=prior, 

702 make_posterior_fn=posterior, 

703 kl_weight=1 / self.x_train_shape, 

704 kl_use_exact=False, 

705 activation='sigmoid'), 

706 

707 DenseVariationalCustom(units=tfpl.IndependentNormal.params_size(self._output_shape), 

708 make_prior_fn=prior, 

709 make_posterior_fn=posterior, 

710 kl_use_exact=False, 

711 kl_weight=1 / self.x_train_shape), 

712 

713 # Aleatoric uncertainty 

714 tfpl.IndependentNormal(self._output_shape) 

715 ]) 

716 logging.warning(model.summary()) 

717 self.model = model 

718 

719 def set_compile_options(self): 

720 self.optimizer = tf.keras.optimizers.RMSprop() 

721 self.loss = nll 

722 

723 

724 

725class ProbTestModelMixture(AbstractModelClass): 

726 def __init__(self, input_shape: list, output_shape: list): 

727 super().__init__(input_shape[0], output_shape[0]) 

728 self.initial_lr = 0.001 

729 self.loss = nll 

730 self.divergence_fn = lambda q, p, _: tfd.kl_divergence(q, p) / input_shape[0][0] 

731 self.k_mixed_components = 2 

732 

733 self.set_model() 

734 self.set_compile_options() 

735 self.set_custom_objects(nll=nll) 

736 

737 def set_model(self): 

738 x_input = tf.keras.layers.Input(self._input_shape) 

739 #tf.keras.layers.Conv2D(kernel_size=(3,1), filters=8, 

740 # activation='relu', padding="same"), 

741 x = Convolution2DReparameterizationCustom( 

742 8, (3, 1), padding='same', 

743 kernel_prior_fn=tfpl.default_multivariate_normal_fn, 

744 kernel_posterior_fn=tfpl.default_mean_field_normal_fn(), 

745 kernel_divergence_fn=self.divergence_fn, 

746 bias_prior_fn=tfpl.default_multivariate_normal_fn, 

747 bias_posterior_fn=tfpl.default_mean_field_normal_fn(), 

748 bias_divergence_fn=self.divergence_fn, 

749 activation='relu', 

750 )(x_input) 

751 

752 x = tf.keras.layers.Flatten()(x) 

753