def train_flipout_nll_model(x_train, y_train, domain):
num_batches = x_train.shape[0] / 32
kl_weight = 1.0 / num_batches
inp = Input(shape=(1, ))
x = FlipoutDense(
32,
kl_weight,
activation="relu",
)(inp)
x = FlipoutDense(32, kl_weight, activation="relu")(x)
mean = FlipoutDense(1, kl_weight, activation="linear")(x)
var = FlipoutDense(1, kl_weight, activation="softplus")(x)
train_model = Model(inp, mean)
pred_model = Model(inp, [mean, var])
train_model.compile(loss=regression_gaussian_nll_loss(var),
optimizer="adam")
train_model.fit(x_train, y_train, verbose=2, epochs=700)
st_model = TwoHeadStochasticRegressor(pred_model)
pred_mean, std_pred = st_model.predict(domain, num_samples=50)
return pred_mean, std_pred
def model_fn():
inp = Input(shape=(1, ))
x = Dense(32, activation="relu")(inp)
x = Dense(32, activation="relu")(x)
mean = Dense(1, activation="linear")(x)
var = Dense(1, activation="softplus")(x)
train_model = Model(inp, mean)
pred_model = Model(inp, [mean, var])
train_model.compile(loss=regression_gaussian_nll_loss(var),
optimizer="adam")
return train_model, pred_model
def train_standard_model(x_train, y_train, domain):
inp = Input(shape=(1, ))
x = Dense(32, activation="relu")(inp)
x = Dense(32, activation="relu")(x)
mean = Dense(1, activation="linear")(x)
var = Dense(1, activation="softplus")(x)
train_model = Model(inp, mean)
pred_model = Model(inp, [mean, var])
train_model.compile(loss=regression_gaussian_nll_loss(var),
optimizer="adam")
train_model.fit(x_train, y_train, verbose=2, epochs=100)
mean_pred, var_pred = pred_model.predict(domain)
std_pred = np.sqrt(var_pred)
return mean_pred, std_pred
def train_stochastic_model(trunk_model,
x_train,
y_train,
domain,
epochs=200,
dense_layer=Dense,
**kwargs):
inp = Input(shape=(1, ))
x = trunk_model(inp)
mean = dense_layer(1, activation="linear", **kwargs)(x)
var = dense_layer(1, activation="softplus", **kwargs)(x)
train_model = Model(inp, mean)
pred_model = Model(inp, [mean, var])
train_model.compile(loss=regression_gaussian_nll_loss(var),
optimizer="adam")
train_model.fit(x_train, y_train, verbose=2, epochs=epochs)
st_model = TwoHeadStochasticRegressor(pred_model)
pred_mean, pred_ale, pred_epi = st_model.predict(
domain, num_samples=NUM_SAMPLES, disentangle_uncertainty=True)
return pred_mean, pred_ale, pred_epi