├── LICENSE ├── README.md ├── base.py ├── batch_norm.py ├── config.json ├── eval.py ├── gmm.py ├── gp.py ├── infer_model.py ├── lets_start.py ├── logger.py ├── model.py ├── ops.py ├── pics └── scheme.png ├── reactions.py ├── realreaction.py ├── rnn.py └── util.py /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2017 Zhenpeng Zhou 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Optimizing Chemical Reactions with Deep Reinforcement Learning 2 | 3 | ![TOC](https://raw.githubusercontent.com/lightingghost/chemopt/master/pics/scheme.png) 4 | 5 | Zhenpeng Zhou, Xiaocheng Li, Richard N. Zare 6 | 7 | The tensorflow implementation of the [paper](http://pubs.acs.org/doi/full/10.1021/acscentsci.7b00492) 8 | 9 | # Abstract 10 | 11 | 12 | Deep reinforcement learning was employed to optimize chemical reactions. Our model iteratively records the results of a chemical reaction and chooses new experimental conditions to improve the reaction outcome. This model outperformed a state-of-the-art blackbox optimization algorithm by using 71% fewer steps on both simulations and real reactions. Furthermore, we introduced an efficient exploration strategy by drawing the reaction conditions from certain probability distributions, which resulted in an improvement on regret from 0.062 to 0.039 compared with a deterministic policy. Combining the efficient exploration policy with accelerated microdroplet reactions, optimal reaction conditions were determined in 30 min for the four reactions considered, and a better understanding of the factors that control microdroplet reactions was reached. Moreover, our model showed a better performance after training on reactions with similar or even dissimilar underlying mechanisms, which demonstrates its learning ability. 13 | 14 | # Getting Started 15 | 16 | edit the hyperparameters in `config.json` and execute 17 | 18 | ```python 19 | python lets_start.py 20 | ``` 21 | 22 | # Use 23 | after training, execute 24 | 25 | ```python 26 | python realreaction.py 27 | ``` 28 | 29 | to access the interactive interface for optimizing a real reaction. 30 | 31 | # Implementation references 32 | 33 | We wish to thank the authors of the following projects for inspiration. 34 | 35 | - [Learning to Learn by Gradient Descent by Gradient Descent](https://github.com/deepmind/learning-to-learn) 36 | -------------------------------------------------------------------------------- /base.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | class Error(Exception): 4 | pass 5 | 6 | class NotConnectedError(Error): 7 | pass 8 | 9 | class ParentNotBuiltError(Error): 10 | pass 11 | 12 | class IncompatibleShapeError(Error): 13 | pass 14 | 15 | class UnderspecifiedError(Error): 16 | pass 17 | 18 | class NotSupportedError(Error): 19 | pass 20 | -------------------------------------------------------------------------------- /batch_norm.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | import base 4 | 5 | from tensorflow.contrib.layers.python.layers import utils 6 | 7 | class BatchNorm(): 8 | GAMMA = "gamma" 9 | BETA = "beta" 10 | POSSIBLE_INITIALIZER_KEYS = {GAMMA, BETA} 11 | 12 | def __init__(self, reduction_indices=None, offset=True, scale=False, 13 | decay_rate=0.999, eps=1e-3, initializers=None, 14 | use_legacy_moving_second_moment=False, 15 | name="batch_norm"): 16 | 17 | self._reduction_indices = reduction_indices 18 | self._offset = offset 19 | self._scale = scale 20 | self._decay_rate = decay_rate 21 | self._eps = eps 22 | self._use_legacy_moving_second_moment = use_legacy_moving_second_moment 23 | self._initializers = util.check_initializers( 24 | initializers, self.POSSIBLE_INITIALIZER_KEYS) 25 | 26 | def _set_default_initializer(self, var_name): 27 | if var_name not in self._initializers: 28 | if var_name == self.GAMMA: 29 | self._initializers[self.GAMMA] = tf.ones_initializer() 30 | elif var_name == self.BETA: 31 | self._initializers[self.BETA] = tf.zeros_initializer 32 | 33 | def _build_statistics_variance(self, input_batch, 34 | reduction_indices, use_batch_stats): 35 | self._moving_mean = tf.get_variable( 36 | "moving_mean", 37 | shape=self._mean_shape, 38 | collections=[tf.GraphKeys.MOVING_AVERAGE_VARIABLES, 39 | tf.GraphKeys.VARIABLES], 40 | initializer=tf.zeros_initializer, 41 | trainable=False) 42 | 43 | self._moving_variance = tf.get_variable( 44 | "moving_variance", 45 | shape=self._mean_shape, 46 | collections=[tf.GraphKeys.MOVING_AVERAGE_VARIABLES, 47 | tf.GraphKeys.VARIABLES], 48 | initializer=tf.ones_initializer(), 49 | trainable=False) 50 | 51 | def build_batch_stats(): 52 | """Builds the batch statistics calculation ops.""" 53 | shift = tf.add(self._moving_mean, 0) 54 | counts, shifted_sum_x, shifted_sum_x2, _ = tf.nn.sufficient_statistics( 55 | input_batch, 56 | reduction_indices, 57 | keep_dims=True, 58 | shift=shift, 59 | name="batch_norm_ss") 60 | 61 | mean, variance = tf.nn.normalize_moments(counts, 62 | shifted_sum_x, 63 | shifted_sum_x2, 64 | shift, 65 | name="normalize_moments") 66 | 67 | return mean, variance 68 | 69 | def build_moving_stats(): 70 | return ( 71 | tf.identity(self._moving_mean), 72 | tf.identity(self._moving_variance),) 73 | 74 | mean, variance = utils.smart_cond( 75 | use_batch_stats, 76 | build_batch_stats, 77 | build_moving_stats, 78 | ) 79 | 80 | return mean, variance 81 | 82 | def _build_statistics_second_moment(self, input_batch, 83 | reduction_indices, use_batch_stats): 84 | self._moving_mean = tf.get_variable( 85 | "moving_mean", 86 | shape=self._mean_shape, 87 | collections=[tf.GraphKeys.MOVING_AVERAGE_VARIABLES, 88 | tf.GraphKeys.VARIABLES], 89 | initializer=tf.zeros_initializer, 90 | trainable=False) 91 | 92 | self._moving_second_moment = tf.get_variable( 93 | "moving_second_moment", 94 | shape=self._mean_shape, 95 | collections=[tf.GraphKeys.MOVING_AVERAGE_VARIABLES, 96 | tf.GraphKeys.VARIABLES], 97 | initializer=tf.ones_initializer(), 98 | trainable=False) 99 | 100 | self._moving_variance = tf.sub(self._moving_second_moment, 101 | tf.square(self._moving_mean), 102 | name="moving_variance") 103 | 104 | def build_batch_stats(): 105 | shift = tf.add(self._moving_mean, 0) 106 | counts, shifted_sum_x, shifted_sum_x2, _ = tf.nn.sufficient_statistics( 107 | input_batch, 108 | reduction_indices, 109 | keep_dims=True, 110 | shift=shift, 111 | name="batch_norm_ss") 112 | 113 | mean, variance = tf.nn.normalize_moments(counts, 114 | shifted_sum_x, 115 | shifted_sum_x2, 116 | shift, 117 | name="normalize_moments") 118 | second_moment = variance + tf.square(mean) 119 | 120 | return mean, variance, second_moment 121 | 122 | def build_moving_stats(): 123 | return ( 124 | tf.identity(self._moving_mean), 125 | tf.identity(self._moving_variance), 126 | tf.identity(self._moving_second_moment), 127 | ) 128 | 129 | mean, variance, second_moment = utils.smart_cond( 130 | use_batch_stats, 131 | build_batch_stats, 132 | build_moving_stats, 133 | ) 134 | 135 | return mean, variance, second_moment 136 | 137 | def _build_update_ops_variance(self, mean, variance, is_training): 138 | def build_update_ops(): 139 | update_mean_op = moving_averages.assign_moving_average( 140 | variable=self._moving_mean, 141 | value=mean, 142 | decay=self._decay_rate, 143 | name="update_moving_mean").op 144 | 145 | update_variance_op = moving_averages.assign_moving_average( 146 | variable=self._moving_variance, 147 | value=variance, 148 | decay=self._decay_rate, 149 | name="update_moving_variance").op 150 | 151 | return update_mean_op, update_variance_op 152 | 153 | def build_no_ops(): 154 | return (tf.no_op(), tf.no_op()) 155 | 156 | # Only make the ops if we know that `is_training=True`, or the 157 | # value of `is_training` is unknown. 158 | is_training_const = utils.constant_value(is_training) 159 | if is_training_const is None or is_training_const: 160 | update_mean_op, update_variance_op = utils.smart_cond( 161 | is_training, 162 | build_update_ops, 163 | build_no_ops, 164 | ) 165 | 166 | # Every new connection creates a new op which adds its contribution 167 | # to the running average when ran. 168 | tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_mean_op) 169 | tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_variance_op) 170 | 171 | def _build_update_ops_second_moment(self, mean, second_moment, is_training): 172 | def build_update_ops(): 173 | update_mean_op = moving_averages.assign_moving_average( 174 | variable=self._moving_mean, 175 | value=mean, 176 | decay=self._decay_rate, 177 | name="update_moving_mean").op 178 | 179 | update_second_moment_op = moving_averages.assign_moving_average( 180 | variable=self._moving_second_moment, 181 | value=second_moment, 182 | decay=self._decay_rate, 183 | name="update_moving_second_moment").op 184 | 185 | return update_mean_op, update_second_moment_op 186 | 187 | def build_no_ops(): 188 | return (tf.no_op(), tf.no_op()) 189 | 190 | is_training_const = utils.constant_value(is_training) 191 | if is_training_const is None or is_training_const: 192 | update_mean_op, update_second_moment_op = utils.smart_cond( 193 | is_training, 194 | build_update_ops, 195 | build_no_ops, 196 | ) 197 | 198 | tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_mean_op) 199 | tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_second_moment_op) 200 | 201 | def _build(self, input_batch, is_training=True, test_local_stats=True): 202 | 203 | input_shape = input_batch.get_shape() 204 | 205 | if self._reduction_indices is not None: 206 | if len(self._reduction_indices) > len(input_shape): 207 | raise base.IncompatibleShapeError( 208 | "Too many reduction indices specified.") 209 | 210 | if max(self._reduction_indices) >= len(input_shape): 211 | raise base.IncompatibleShapeError( 212 | "Reduction index too large for input shape.") 213 | 214 | if min(self._reduction_indices) < 0: 215 | raise base.IncompatibleShapeError( 216 | "Reduction indeces must be non-negative.") 217 | 218 | reduction_indices = self._reduction_indices 219 | else: 220 | reduction_indices = range(len(input_shape))[:-1] 221 | 222 | if input_batch.dtype == tf.float16: 223 | raise base.NotSupportedError( 224 | "BatchNorm does not support `tf.float16`, insufficient " 225 | "precision for calculating sufficient statistics.") 226 | 227 | self._mean_shape = input_batch.get_shape().as_list() 228 | for index in reduction_indices: 229 | self._mean_shape[index] = 1 230 | 231 | use_batch_stats = is_training | test_local_stats 232 | 233 | # Use the legacy moving second moment if the flag is set. 234 | if self._use_legacy_moving_second_moment: 235 | tf.logging.warning( 236 | "nn.BatchNorm `use_legacy_second_moment=True` is deprecated.") 237 | 238 | mean, variance, second_moment = self._build_statistics_second_moment( 239 | input_batch, 240 | reduction_indices, 241 | use_batch_stats) 242 | 243 | self._build_update_ops_second_moment(mean, second_moment, is_training) 244 | else: 245 | mean, variance = self._build_statistics_variance( 246 | input_batch, 247 | reduction_indices, 248 | use_batch_stats) 249 | 250 | self._build_update_ops_variance(mean, variance, is_training) 251 | 252 | # Set up optional scale and offset factors. 253 | if self._offset: 254 | self._set_default_initializer(self.BETA) 255 | self._beta = tf.get_variable( 256 | self.BETA, 257 | shape=self._mean_shape, 258 | initializer=self._initializers[self.BETA]) 259 | else: 260 | self._beta = None 261 | 262 | if self._scale: 263 | self._set_default_initializer(self.GAMMA) 264 | self._gamma = tf.get_variable( 265 | self.GAMMA, 266 | shape=self._mean_shape, 267 | initializer=self._initializers[self.GAMMA]) 268 | else: 269 | self._gamma = None 270 | 271 | out = tf.nn.batch_normalization( 272 | input_batch, 273 | mean, 274 | variance, 275 | self._beta, 276 | self._gamma, 277 | self._eps, 278 | name="batch_norm") 279 | 280 | return out 281 | 282 | @property 283 | def moving_mean(self): 284 | self._ensure_is_connected() 285 | return self._moving_mean 286 | 287 | @property 288 | def moving_second_moment(self): 289 | self._ensure_is_connected() 290 | return self._moving_second_moment 291 | 292 | @property 293 | def moving_variance(self): 294 | self._ensure_is_connected() 295 | return self._moving_variance 296 | 297 | @property 298 | def beta(self): 299 | self._ensure_is_connected() 300 | 301 | if self._beta is None: 302 | raise base.Error( 303 | "Batch normalization doesn't have an offset, so no beta") 304 | else: 305 | return self._beta 306 | 307 | @property 308 | def gamma(self): 309 | self._ensure_is_connected() 310 | 311 | if self._gamma is None: 312 | raise base.Error( 313 | "Batch normalization doesn't have a scale, so no gamma") 314 | else: 315 | return self._gamma 316 | -------------------------------------------------------------------------------- /config.json: -------------------------------------------------------------------------------- 1 | { 2 | "log_path": "./log/", 3 | "save_path": "./ckpt/8_dim_quad/h128_s50_srnn/", 4 | "batch_size": 128, 5 | "hidden_size": 128, 6 | "num_layers": 2, 7 | "batch_norm": false, 8 | "reuse": false, 9 | "num_epochs": 50000, 10 | "log_period": 100, 11 | "evaluation_period": 1000, 12 | "evaluation_epochs": 20, 13 | "reaction_type": "quad", 14 | "norm_cov": 0.3, 15 | "constraints": false, 16 | "num_params": 3, 17 | "instrument_error": null, 18 | "num_steps": 50, 19 | "unroll_length": 50, 20 | "learning_rate": 0.001, 21 | "optimizer": "Adam", 22 | "loss_type": "oi", 23 | "discount_factor": 0.97, 24 | "opt_direction": "min", 25 | "policy": "srnn", 26 | "trainable_init": true 27 | 28 | } 29 | -------------------------------------------------------------------------------- /eval.py: -------------------------------------------------------------------------------- 1 | import util 2 | import logging 3 | import json 4 | import os 5 | os.environ['TF_CPP_MIN_LOG_LEVEL']='3' 6 | import tensorflow as tf 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | 10 | from model import Optimizer 11 | from rnn import MultiInputLSTM 12 | from logger import get_handlers 13 | from collections import namedtuple 14 | 15 | logging.basicConfig(level=logging.INFO, handlers=get_handlers(False)) 16 | logger = logging.getLogger() 17 | 18 | 19 | def main(): 20 | 21 | 22 | config_file = open('./config.json') 23 | config = json.load(config_file, 24 | object_hook=lambda d:namedtuple('x', d.keys())(*d.values())) 25 | num_unrolls = config.num_steps // config.unroll_length 26 | with tf.Session() as sess: 27 | model = util.load_model(sess, config, logger) 28 | all_y = [] 29 | for i in range(10): 30 | print(i) 31 | _, loss, reset, fx_array, x_array = model.step() 32 | cost, others = util.run_epoch(sess, loss, [fx_array, x_array], 33 | reset, num_unrolls) 34 | Y, X = others 35 | all_y.append(Y) 36 | 37 | all_y = np.hstack(all_y) 38 | np.save('srnn.npy', all_y) 39 | plt.figure(1) 40 | y_mean = np.mean(all_y, axis=1) 41 | plt.plot(y_mean) 42 | print(min(y_mean)) 43 | plt.show() 44 | 45 | 46 | if __name__ == '__main__': 47 | main() 48 | -------------------------------------------------------------------------------- /gmm.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.stats import multivariate_normal as normal 3 | import matplotlib.pyplot as plt 4 | from matplotlib import cm 5 | from itertools import product 6 | from mpl_toolkits.mplot3d import Axes3D 7 | 8 | import tensorflow as tf 9 | from reactions import GMM as tf_GMM 10 | 11 | class GMM: 12 | def __init__(self, n=6, ndim=3, cov=0.15, record=False): 13 | self.n = n 14 | self.ndim = ndim 15 | self.record = record 16 | self.cov = cov 17 | self.refresh() 18 | 19 | def refresh(self): 20 | self.m = [np.random.rand(self.ndim) for _ in range(self.n)] 21 | self.cov = [np.random.normal(self.cov, self.cov/5, size=self.ndim) 22 | for _ in range(self.n)] 23 | self.param = np.random.normal(loc=0, scale=0.2, size=self.n) 24 | self.param /= np.sum(np.abs(self.param)) 25 | if self.record: 26 | self.history = {'x':[], 'y':[]} 27 | 28 | self.cst = (2 * 3.14159) ** (- self.ndim / 2) 29 | modes = np.array([1/np.prod(cov) for cov in self.cov]) 30 | modes = modes * self.param 31 | self.tops = np.max(modes) 32 | self.bots = np.min(modes) 33 | 34 | def __call__(self, x): 35 | y = [normal.pdf(x, self.m[i], self.cov[i]) for i in range(self.n)] 36 | 37 | fx = np.asscalar( 38 | np.dot( 39 | self.param.reshape((1, -1)), 40 | np.array(y).reshape((-1, 1)))/self.n) 41 | result = (fx / self.cst - self.bots) / (self.tops - self.bots) 42 | if self.record: 43 | self.history['x'].append(x) 44 | self.history['y'].append(result) 45 | return result 46 | 47 | 48 | 49 | def test_1d(): 50 | gmm = GMM(ndim=1) 51 | x = np.arange(0, 1, 0.01) 52 | y = [gmm(i) for i in x] 53 | plt.figure(1) 54 | plt.plot(x, y) 55 | plt.show() 56 | 57 | def test_2d(): 58 | gmm = GMM(ndim=2) 59 | xr = list(np.arange(0, 1, 0.02)) 60 | X = np.array(list(product(xr, repeat=2))) 61 | Y = [gmm(i) for i in X] 62 | fig = plt.figure(1) 63 | ax = fig.gca(projection='3d') 64 | ax.plot_trisurf(X[:, 0], X[:, 1], Y) 65 | fig.show() 66 | plt.show() 67 | 68 | def test_tf(): 69 | xr = list(np.arange(0, 1, 0.02)) 70 | X = np.array(list(product(xr, repeat=2))) 71 | Y = [] 72 | with tf.Session() as sess: 73 | gmm = tf_GMM(batch_size=1, ncoef=6, num_dims=2, cov=0.5) 74 | y = gmm(tf.placeholder(tf.float32, shape=[1, 2], name='x')) 75 | sess.run(tf.global_variables_initializer()) 76 | for x in X: 77 | Y.append(sess.run(y, feed_dict={'x:0':x.reshape((1, 2))})) 78 | 79 | cmap = cm.get_cmap('rainbow') 80 | fig = plt.figure(1) 81 | ax = fig.gca(projection='3d') 82 | ax.plot_trisurf(X[:, 0], X[:, 1], np.squeeze(Y), 83 | linewidth=0.0, antialiased=True, 84 | cmap=cmap) 85 | fig.show() 86 | plt.show() 87 | 88 | if __name__ == '__main__': 89 | test_tf() 90 | -------------------------------------------------------------------------------- /gp.py: -------------------------------------------------------------------------------- 1 | import GPy 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | from itertools import product 5 | from reactions import QuadraticEval 6 | 7 | def gp(prior_X, prior_Y, variance=0.1, lengthscale=0.1, X=None, nsamples=1): 8 | nd = prior_X.shape[1] 9 | kernel = GPy.kern.RBF(input_dim=nd, variance=variance, lengthscale=lengthscale) 10 | m = GPy.models.GPRegression(prior_X, prior_Y, kernel) 11 | if X is None: 12 | X = np.arange(0, 1, 0.01).reshape((-1, 1)) 13 | Y = m.posterior_samples_f(X, size=nsamples) 14 | return Y 15 | 16 | def plot_1d(X, Y): 17 | plt.figure(1) 18 | plt.plot(X, Y) 19 | plt.show() 20 | 21 | def test_1d(): 22 | prior_X = np.array([[0], [1]]) 23 | prior_Y = np.array([[10], [10]]) 24 | X = np.arange(0, 1, 0.01) 25 | Y = gp(prior_X, prior_Y, nsamples=1) 26 | plot_1d(X, Y) 27 | 28 | def plot_2d(X, Y, Z): 29 | fig = plt.figure(1) 30 | ax = fig.gca(projection='3d') 31 | ax.plot_trisurf(X, Y, np.squeeze(Z)) 32 | plt.show() 33 | 34 | def test_2d(): 35 | prior_X = np.array([[0, 0]]) 36 | prior_Y = np.array([[0]]) 37 | xr = list(np.arange(0, 1, 0.02)) 38 | X = np.array(list(product(xr, repeat=2))) 39 | Y = gp(prior_X, prior_Y, X = X, nsamples=1, lengthscale=0.3) 40 | plot_2d(X[:, 0], X[:, 1], Y) 41 | 42 | 43 | class SquaredDistanceKernel(): 44 | def __init__(self, param=0.1): 45 | self.param = param 46 | 47 | def __call__(self, a, b): 48 | sq_dist = np.sum(a ** 2, 1).reshape(-1, 1) + np.sum(b ** 2, 1) - 2 * np.dot(a, b.T) 49 | return np.exp(- sq_dist / self.param / 2) 50 | 51 | class GaussianProcess(object): 52 | def __init__(self, kernel=SquaredDistanceKernel(), noise=0.0): 53 | self.kernel = kernel 54 | self.noise = noise 55 | self.X = None 56 | self.Y = None 57 | 58 | def prior(self, x, y): 59 | self.X = x 60 | self.Y = y 61 | 62 | def predict(self, x): 63 | k2 = self.kernel(x, x) 64 | self.cov = None 65 | if self.X and self.Y: 66 | self.cov = self.kernel(X) 67 | if self.cov is None: 68 | mu = np.zeros(x.shape) 69 | cov_posterior = k2 + (self.noise * np.eye(k2.shape[0])) 70 | 71 | else: 72 | l = np.linalg.cholesky(self.cov + self.noise * np.eye(self.cov.shape[0])) 73 | k = self.kernel(X, x) 74 | ldk = np.linalg.solve(l, k) 75 | 76 | mu = np.dot(ldk.T, np.linalg.solve(1, self.Y)) 77 | cov_posterior = k2 + self.noise * np.eye(k2.shape[0]) - np.dot(ldk.T, ldk) 78 | 79 | return mu, cov_posterior 80 | 81 | def t1d(): 82 | gp = GaussianProcess() 83 | np.random.seed(1) 84 | x = np.arange(0, 1, 0.01).reshape(-1, 1) 85 | mu, cov = gp.predict(x) 86 | y = np.random.multivariate_normal(np.squeeze(mu), cov) 87 | plot_1d(np.squeeze(x), np.squeeze(y)) 88 | 89 | def t2d(): 90 | gp = GaussianProcess() 91 | np.random.seed(1) 92 | xr = list(np.arange(0, 1, 0.02)) 93 | x = np.array(list(product(xr, repeat=2))) 94 | mu, cov = gp.predict(np.array([[0.48, 0.68], [0, 0]])) 95 | import pdb; pdb.set_trace() 96 | y = np.random.multivariate_normal(np.zeros(2), cov) 97 | plot_2d(x[:, 0], x[:, 1], y) 98 | 99 | class GPOpt: 100 | def __init__(self, ndim, prange=[]): 101 | self.ndim = ndim 102 | self.prange = prange 103 | self.X = [] 104 | self.y = [] 105 | xr = list(np.arange(0, 1, 0.1)) 106 | self.x = np.array(list(product(xr, repeat=ndim))) 107 | 108 | def update(self, X, y): 109 | normalized_X = [0] * self.ndim 110 | for i in range(self.ndim): 111 | a, b = self.prange[i] 112 | normalized_X[i] = (X[i] - a) / (b - a) 113 | self.X.append(normalized_X) 114 | self.y.append(y) 115 | 116 | def next(self): 117 | if len(self.X) == 0: 118 | x = np.random.rand(3) 119 | else: 120 | kernel = GPy.kern.RBF(input_dim=self.ndim, 121 | variance=1, lengthscale=1) 122 | X = np.array(self.X) 123 | y = np.array(self.y).reshape((-1, 1)) 124 | m = GPy.models.GPRegression(X, y, kernel) 125 | y_pred = m.posterior_samples_f(self.x, size=1) 126 | x = self.x[np.argmax(y_pred)] 127 | real_x = [0] * self.ndim 128 | for i in range(self.ndim): 129 | a, b = self.prange[i] 130 | real_x[i] = x[i] * (b - a) + a 131 | return real_x 132 | 133 | def test_gpopt(): 134 | opt = GPOpt(3, prange=[(0, 2), (0, 2), (0, 2)]) 135 | func = QuadraticEval(num_dim=3, random=None, ptype='concave') 136 | 137 | y_array = [] 138 | for i in range(30): 139 | x = opt.next() 140 | y_eval = func(x) 141 | y_array.append(y_eval) 142 | opt.update(x, y_eval) 143 | 144 | plt.figure() 145 | plt.plot(y_array) 146 | plt.show() 147 | 148 | test_gpopt() 149 | 150 | 151 | -------------------------------------------------------------------------------- /infer_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL']='3' 3 | import tensorflow as tf 4 | import numpy as np 5 | import logging 6 | import matplotlib.pyplot as plt 7 | from mpl_toolkits.mplot3d import Axes3D 8 | import json 9 | 10 | import rnn 11 | from reactions import QuadraticEval, ConstraintQuadraticEval, RealReaction 12 | from logger import get_handlers 13 | from collections import namedtuple 14 | 15 | logging.basicConfig(level=logging.INFO, handlers=get_handlers()) 16 | logger = logging.getLogger() 17 | 18 | 19 | class StepOptimizer: 20 | def __init__(self, cell, func, ndim, nsteps, ckpt_path, logger, constraints): 21 | self.logger = logger 22 | self.cell = cell 23 | self.func = func 24 | self.ndim = ndim 25 | self.nsteps = nsteps 26 | self.ckpt_path = ckpt_path 27 | self.constraints = constraints 28 | self.init_state = self.cell.get_initial_state(1, tf.float32) 29 | self.results = self.build_graph() 30 | 31 | self.saver = tf.train.Saver(tf.global_variables()) 32 | 33 | def get_state_shapes(self): 34 | return [(s[0].get_shape().as_list(), s[1].get_shape().as_list()) 35 | for s in self.init_state] 36 | 37 | def step(self, sess, x, y, state): 38 | feed_dict = {'input_x:0':x, 'input_y:0':y} 39 | for i in range(len(self.init_state)): 40 | feed_dict['state_l{0}_c:0'.format(i)] = state[i][0] 41 | feed_dict['state_l{0}_h:0'.format(i)] = state[i][1] 42 | new_x, new_state = sess.run(self.results, feed_dict=feed_dict) 43 | return new_x, new_state 44 | 45 | def build_graph(self): 46 | x = tf.placeholder(tf.float32, shape=[1, self.ndim], name='input_x') 47 | y = tf.placeholder(tf.float32, shape=[1, 1], name='input_y') 48 | state = [] 49 | for i in range(len(self.init_state)): 50 | state.append((tf.placeholder( 51 | tf.float32, shape=self.init_state[i][0].get_shape(), 52 | name='state_l{0}_c'.format(i)), 53 | tf.placeholder( 54 | tf.float32, shape=self.init_state[i][1].get_shape(), 55 | name='state_l{0}_h'.format(i)))) 56 | 57 | with tf.name_scope('opt_cell'): 58 | new_x, new_state = self.cell(x, y, state) 59 | if self.constraints: 60 | new_x = tf.clip_by_value(new_x, 0.01, 0.99) 61 | return new_x, new_state 62 | 63 | def load(self, sess, ckpt_path): 64 | ckpt = tf.train.get_checkpoint_state(ckpt_path) 65 | if ckpt and ckpt.model_checkpoint_path: 66 | logger.info('Reading model parameters from {}.'.format( 67 | ckpt.model_checkpoint_path)) 68 | self.saver.restore(sess, ckpt.model_checkpoint_path) 69 | else: 70 | raise FileNotFoundError('No checkpoint available') 71 | 72 | def get_init(self): 73 | x = np.random.normal(loc=0.5, scale=0.2, size=(1, 3)) 74 | x = np.maximum(np.minimum(x, 0.9), 0.1) 75 | y = np.array(self.func(x)).reshape(1, 1) 76 | init_state = [(np.zeros(s[0]), np.zeros(s[1])) 77 | for s in self.get_state_shapes()] 78 | return x, y, init_state 79 | 80 | def run(self): 81 | with tf.Session() as sess: 82 | self.load(sess, self.ckpt_path) 83 | x, y, state = self.get_init() 84 | x_array = np.zeros((self.nsteps + 1, self.ndim)) 85 | y_array = np.zeros((self.nsteps + 1, 1)) 86 | x_array[0, :] = x 87 | y_array[0] = y 88 | for i in range(self.nsteps): 89 | x, state = self.step(sess, x, y, state) 90 | y = np.array(self.func(x)).reshape(1, 1) 91 | x_array[i+1, :] = x 92 | y_array[i+1] = y 93 | 94 | return x_array, y_array 95 | 96 | def main(): 97 | config_file = open('./config.json') 98 | config = json.load(config_file, 99 | object_hook=lambda d:namedtuple('x', d.keys())(*d.values())) 100 | 101 | if config.opt_direction is 'max': 102 | problem_type = 'concave' 103 | else: 104 | problem_type = 'convex' 105 | 106 | if config.constraints: 107 | func = ConstraintQuadraticEval(num_dim=config.num_params, 108 | random=config.instrument_error, 109 | ptype=problem_type) 110 | else: 111 | func = QuadraticEval(num_dim=config.num_params, 112 | random=config.instrument_error, 113 | ptype=problem_type) 114 | 115 | if config.policy == 'srnn': 116 | cell = rnn.StochasticRNNCell(cell=rnn.LSTM, 117 | kwargs= 118 | {'hidden_size':config.hidden_size, 119 | 'use_batch_norm_h':config.batch_norm, 120 | 'use_batch_norm_x':config.batch_norm, 121 | 'use_batch_norm_c':config.batch_norm,}, 122 | nlayers=config.num_layers, 123 | reuse=config.reuse) 124 | if config.policy == 'rnn': 125 | cell = rnn.MultiInputRNNCell(cell=rnn.LSTM, 126 | kwargs= 127 | {'hidden_size':config.hidden_size, 128 | 'use_batch_norm_h':config.batch_norm, 129 | 'use_batch_norm_x':config.batch_norm, 130 | 'use_batch_norm_c':config.batch_norm,}, 131 | nlayers=config.num_layers, 132 | reuse=config.reuse) 133 | 134 | optimizer = StepOptimizer(cell=cell, func=func, ndim=config.num_params, 135 | nsteps=config.num_steps, 136 | ckpt_path=config.save_path, logger=logger, 137 | constraints=config.constraints) 138 | x_array, y_array = optimizer.run() 139 | 140 | # np.savetxt('./scratch/nn_y.csv', y_array, delimiter=',') 141 | # np.save('./scratch/nn_x.npy', y_array) 142 | plt.figure(1) 143 | plt.plot(y_array) 144 | plt.show() 145 | fig2 = plt.figure(2) 146 | ax2 = fig2.add_subplot(111, projection='3d') 147 | ax2.plot(x_array[:, 0], x_array[:, 1], x_array[:, 2]) 148 | fig2.show() 149 | plt.show() 150 | 151 | if __name__ == '__main__': 152 | main() 153 | -------------------------------------------------------------------------------- /lets_start.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL']='3' 3 | import tensorflow as tf 4 | import util 5 | import numpy as np 6 | import logging 7 | import json 8 | 9 | from model import Optimizer 10 | from rnn import MultiInputLSTM 11 | from logger import get_handlers 12 | from collections import namedtuple 13 | 14 | logging.basicConfig(level=logging.INFO, handlers=get_handlers()) 15 | logger = logging.getLogger() 16 | 17 | def main(): 18 | config_file = open('./config.json') 19 | config = json.load(config_file, 20 | object_hook=lambda d:namedtuple('x', d.keys())(*d.values())) 21 | config_file.seek(0) 22 | logger.info(str(json.load(config_file))) 23 | config_file.close() 24 | num_unrolls = config.num_steps // config.unroll_length 25 | with tf.Session() as sess: 26 | # tf.get_default_graph().finalize() 27 | model = util.create_model(sess, config, logger) 28 | step, loss, reset, fx_array, x_array = model.step() 29 | 30 | best_cost = [float('inf')] * 3 31 | epoch_cost = 0 32 | total_cost = 0 33 | 34 | for e in range(config.num_epochs): 35 | cost, _ = util.run_epoch(sess, loss, [step], reset, num_unrolls) 36 | epoch_cost += cost 37 | total_cost += cost 38 | 39 | if (e + 1) % config.log_period == 0: 40 | lm_e = epoch_cost / config.log_period 41 | logger.info('Epoch {}, Mean Error: {:.3f}'.format(e, lm_e)) 42 | epoch_cost = 0 43 | 44 | if (e + 1) % config.evaluation_period == 0: 45 | elm_e = total_cost / config.evaluation_period 46 | logger.info('Current {} epochs, Mean Error: {:.3f}'.format(config.evaluation_period, elm_e)) 47 | 48 | mbc = max(best_cost) 49 | if config.save_path is not None and total_cost < mbc: 50 | best_cost.remove(mbc) 51 | best_cost.append(total_cost) 52 | logger.info('Save current model ...') 53 | model.saver.save(sess, config.save_path, global_step=e) 54 | 55 | total_cost = 0 56 | 57 | 58 | if __name__ == '__main__': 59 | main() 60 | -------------------------------------------------------------------------------- /logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import datetime 3 | import time 4 | import os 5 | import sys 6 | 7 | def get_handlers(log_file=True, log_stdout=True): 8 | handlers = [] 9 | ts = time.time() 10 | date = datetime.datetime.fromtimestamp(ts).strftime('%m-%d-%y') 11 | if not os.path.isdir(os.path.join('.', 'log')): 12 | os.mkdir(os.path.join('.', 'log')) 13 | log_filepath = os.path.join('.', 'log', date + '.log') 14 | 15 | log_fmt = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s') 16 | 17 | if log_file: 18 | file_handler = logging.FileHandler(log_filepath) 19 | file_handler.setFormatter(log_fmt) 20 | handlers.append(file_handler) 21 | if log_stdout: 22 | console_handler = logging.StreamHandler(sys.stdout) 23 | console_handler.setFormatter(log_fmt) 24 | handlers.append(console_handler) 25 | 26 | return handlers 27 | 28 | def set_logger(name='root', log_file=True, log_stdout=True): 29 | ts = time.time() 30 | date = datetime.datetime.fromtimestamp(ts).strftime('%m-%d-%y') 31 | if not os.path.isdir(os.path.join('.', 'log')): 32 | os.mkdir(os.path.join('.', 'log')) 33 | log_filepath = os.path.join('.', 'log', date + '.log') 34 | 35 | log_fmt = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s') 36 | root_logger = logging.getLogger(name) 37 | if log_file: 38 | file_handler = logging.FileHandler(log_filepath) 39 | file_handler.setFormatter(log_fmt) 40 | root_logger.addHandler(file_handler) 41 | if log_stdout: 42 | console_handler = logging.StreamHandler(sys.stdout) 43 | console_handler.setFormatter(log_fmt) 44 | root_logger.addHandler(console_handler) 45 | 46 | 47 | def get_logger(name='root', log_file=True, log_stdout=True): 48 | set_logger(name, log_file, log_stdout) 49 | return logging.getLogger(name) 50 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import ops 3 | import pdb 4 | 5 | from tensorflow.python.util import nest 6 | 7 | class Optimizer: 8 | def __init__(self, cell, logger, func, ndim, batch_size, unroll_len, 9 | lr=0.01, loss_type='naive', optimizer='Adam', trainable_init=False, 10 | direction='max', constraints=False, discount_factor=1.0): 11 | self.batch_size = batch_size 12 | self.constraints = constraints 13 | self.logger = logger 14 | self.cell = cell 15 | self.trainable_init = trainable_init 16 | self.df = self.make_discount(discount_factor, unroll_len) 17 | self.make_loss(func, ndim, batch_size, unroll_len) 18 | loss_func = self.get_loss_func(loss_type, direction) 19 | self.loss = loss_func(self.fx_array) 20 | optimizer = getattr(tf.train, optimizer + 'Optimizer')(lr) 21 | gvs = optimizer.compute_gradients(self.loss) 22 | capped_gvs = [(tf.clip_by_value(grad, -0.1, 0.1), var) for grad, var in gvs] 23 | self.opt = optimizer.apply_gradients(capped_gvs) 24 | 25 | # self.opt = optimizer.minimize(self.loss) 26 | self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=3) 27 | logger.info('model variable:') 28 | logger.info(str([var.name for var in tf.global_variables()])) 29 | logger.info('trainable variables:') 30 | logger.info(str([var.name for var in tf.trainable_variables()])) 31 | self.fx_array = self.fx_array.stack() 32 | self.x_array = self.x_array.stack() 33 | 34 | 35 | def make_discount(self, gamma, unroll_len): 36 | df = [(gamma ** (unroll_len - i)) for i in range(unroll_len + 1)] 37 | return tf.constant(df, shape=[unroll_len + 1, 1], dtype=tf.float32) 38 | 39 | 40 | def make_loss(self, func, ndim, batch_size, unroll_len): 41 | self.unroll_len = unroll_len 42 | x = tf.get_variable('x', shape=[batch_size, ndim], 43 | initializer=tf.truncated_normal_initializer(mean=0.5, stddev=0.2), 44 | trainable=self.trainable_init) 45 | constants = func.get_parameters() 46 | state = self.cell.get_initial_state(batch_size, tf.float32) 47 | self.fx_array = tf.TensorArray(tf.float32, 48 | size=unroll_len+1, clear_after_read=False) 49 | self.x_array = tf.TensorArray(tf.float32, 50 | size=unroll_len+1, clear_after_read=False) 51 | 52 | def step(t, x, state, fx_array, x_array): 53 | with tf.name_scope('fx'): 54 | fx = func(x) 55 | fx_array = fx_array.write(t, fx) 56 | x_array = x_array.write(t, x) 57 | with tf.name_scope('opt_cell'): 58 | new_x, new_state = self.cell(x, fx, state) 59 | if self.constraints: 60 | new_x = tf.clip_by_value(new_x, 0.01, 0.99) 61 | 62 | with tf.name_scope('t_next'): 63 | t_next = t + 1 64 | 65 | return t_next, new_x, new_state, fx_array, x_array 66 | 67 | _, x_final, s_final, self.fx_array, self.x_array = tf.while_loop( 68 | cond=lambda t, *_: t < unroll_len, 69 | body=step, loop_vars=(0, x, state, self.fx_array, self.x_array), 70 | parallel_iterations=1, 71 | swap_memory=True 72 | ) 73 | 74 | with tf.name_scope('fx'): 75 | fx_final = func(x_final) 76 | self.fx_array = self.fx_array.write(unroll_len, fx_final) 77 | self.x_array = self.x_array.write(unroll_len, x) 78 | 79 | # Reset the state; should be called at the beginning of an epoch. 80 | with tf.name_scope('reset'): 81 | 82 | variables = [x,] + constants 83 | # Empty array as part of the reset process. 84 | self.reset = [tf.variables_initializer(variables), 85 | self.fx_array.close(), self.x_array.close()] 86 | 87 | return self.fx_array, self.x_array 88 | 89 | def get_loss_func(self, loss_type='naive', direction='max'): 90 | def loss_func(fx): 91 | if loss_type == 'naive': 92 | loss = tf.reduce_sum( 93 | tf.matmul(tf.reshape(fx.stack(), [self.batch_size, -1]), 94 | self.df, name='loss')) 95 | elif loss_type == 'oi' and direction == 'max': 96 | loss = tf.reduce_sum( 97 | [fx.read(i) - tf.reduce_max( 98 | fx.gather(list(range(i))), axis=0) 99 | for i in range(1, self.unroll_len + 1)], 100 | name='loss') 101 | elif loss_type == 'oi' and direction == 'min': 102 | loss = tf.reduce_sum( 103 | [fx.read(i) - tf.reduce_min( 104 | fx.gather(list(range(i))), axis=0) 105 | for i in range(1, self.unroll_len + 1)], 106 | name='loss') 107 | if direction == 'max': 108 | loss = - loss 109 | return loss / self.batch_size 110 | return loss_func 111 | 112 | def step(self): 113 | return self.opt, self.loss, self.reset, self.fx_array, self.x_array 114 | -------------------------------------------------------------------------------- /ops.py: -------------------------------------------------------------------------------- 1 | import collections 2 | 3 | import tensorflow as tf 4 | import mock 5 | 6 | def wrap_variable_creation(func, custom_getter): 7 | """Provides a custom getter for all variable creations.""" 8 | original_get_variable = tf.get_variable 9 | def custom_get_variable(*args, **kwargs): 10 | if hasattr(kwargs, 'custom_getter'): 11 | raise AttributeError('Custom getters are not supported for ' 12 | 'optimizee variables.') 13 | return original_get_variable(*args, custom_getter=custom_getter, **kwargs) 14 | # Mock the get_variable method. 15 | with mock.patch("tensorflow.get_variable", custom_get_variable): 16 | return func() 17 | 18 | def get_variables(func): 19 | """Calls func, returning any variables created, but ignoring its return value. 20 | 21 | Args: 22 | func: Function to be called. 23 | 24 | Returns: 25 | A tuple (variables, constants) where the first element is a list of 26 | trainable variables and the second is the non-trainable variables. 27 | """ 28 | variables = [] 29 | constants = [] 30 | 31 | def custom_getter(getter, name, **kwargs): 32 | trainable = kwargs['trainable'] 33 | kwargs['trainable'] = False 34 | variable = getter(name, **kwargs) 35 | if trainable: 36 | variables.append(variable) 37 | else: 38 | constants.append(variable) 39 | return variable 40 | 41 | with tf.name_scope("unused_graph"): 42 | wrap_variable_creation(func, custom_getter) 43 | 44 | return variables, constants 45 | 46 | def run_with_custom_variables(func, variable): 47 | """Calls func and replaces any trainable variables. 48 | 49 | This returns the output of func, but whenever `get_variable` is called it 50 | will replace any trainable variables with the tensors in `variables`, in 51 | the same order. Non-trainable variables will re-use any variables already 52 | created. 53 | 54 | Args: 55 | func: Function to be called. 56 | variables: A list of tensors replacing the trainable variables. 57 | 58 | Returns: 59 | The return value of func is returned. 60 | """ 61 | variables = collections.deque(variables) 62 | 63 | def custom_getter(getter, name, **kwargs): 64 | if kwargs["trainable"]: 65 | return variables.popleft() 66 | else: 67 | kwargs["reuse"] = True 68 | return getter(name, **kwargs) 69 | 70 | return wrap_variable_creation(func, custom_getter) 71 | -------------------------------------------------------------------------------- /pics/scheme.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightingghost/chemopt/926440bd2c976a8e259bbb1d2f21399b93ebd210/pics/scheme.png -------------------------------------------------------------------------------- /reactions.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL']='3' 3 | import tensorflow as tf 4 | import numpy as np 5 | 6 | class ConstraintQuadratic: 7 | """Quadratic problem: f(x) = ||Wx - y||.""" 8 | def __init__(self, batch_size=128, num_dims=3, ptype='convex', 9 | random=0.05, dtype=tf.float32): 10 | self.ptype = ptype 11 | self.w = tf.get_variable('w', shape=[batch_size, num_dims, num_dims], 12 | dtype=dtype, initializer=tf.random_normal_initializer(), 13 | trainable=False) 14 | 15 | self.a = tf.get_variable('y', shape=[batch_size, num_dims], 16 | dtype=dtype, initializer=tf.random_uniform_initializer(minval=0.01, maxval=0.99), 17 | trainable=False) 18 | 19 | self.y = tf.squeeze(tf.matmul(self.w, tf.expand_dims(self.a, -1))) 20 | 21 | self.normalizer = tf.maximum( 22 | self._func(tf.zeros([batch_size, num_dims])), 23 | self._func(tf.ones([batch_size, num_dims]))) 24 | 25 | if random is not None: 26 | self.e = tf.random_normal(shape=[batch_size,], stddev=random, 27 | dtype=dtype, name='e') 28 | else: 29 | self.e = 0.0 30 | 31 | def get_parameters(self): 32 | return [self.w, self.a] 33 | 34 | def _func(self, var): 35 | product = tf.squeeze(tf.matmul(self.w, tf.expand_dims(var, -1))) 36 | norm = tf.reduce_sum((product - self.y) ** 2, 1) 37 | return norm 38 | 39 | def _barrier(self, var): 40 | return -tf.reduce_sum(tf.log(var) + tf.log(1 - var), 1) / 1e10 41 | 42 | def __call__(self, x): 43 | ''' 44 | x = tf.get_variable('x', shape=[batch_size, num_dims], 45 | dtype=dtype, initializer=tf.random_normal_initializer(stddev=stdev)) 46 | ''' 47 | res = (self._func(x) / self.normalizer + self.e + self._barrier(x)) 48 | if self.ptype == 'concave': 49 | res = 1 - res 50 | return res 51 | 52 | class GMM: 53 | def __init__(self, batch_size=128, ncoef=6, num_dims=3, random=None, 54 | cov=0.1, dtype=tf.float32): 55 | self.ncoef = ncoef 56 | self.num_dim = num_dims 57 | self.batch_size = batch_size 58 | self.dtype = dtype 59 | with tf.variable_scope('func_gmm'): 60 | self.m = [tf.get_variable('mu_{}'.format(i), shape=[batch_size, num_dims], 61 | dtype=dtype, 62 | initializer=tf.random_uniform_initializer(minval=0.01, maxval=0.99), 63 | trainable=False) 64 | for i in range(ncoef)] 65 | 66 | self.cov = [tf.get_variable('cov_{}'.format(i), shape=[batch_size, num_dims], 67 | dtype=dtype, 68 | initializer=tf.truncated_normal_initializer( 69 | mean=cov, stddev=cov/5), 70 | trainable=False) 71 | for i in range(ncoef)] 72 | 73 | self.coef = tf.get_variable('coef', shape=[ncoef, 1], dtype=dtype, 74 | initializer=tf.random_normal_initializer(stddev=0.2), 75 | trainable=False) 76 | 77 | self.random = random 78 | # if random is not None: 79 | # self.e = tf.random_normal(shape=[batch_size, ], stddev=random, 80 | # dtype=dtype, name='error') 81 | # else: 82 | # self.e = 0.0 83 | 84 | self.cst = (2 * 3.14159) ** (- self.num_dim / 2) 85 | modes = tf.concat([(1 / tf.reduce_prod(cov, axis=1, keep_dims=True)) 86 | for cov in self.cov], axis=1) * tf.transpose(self.coef) 87 | self.tops = tf.reduce_max(modes, axis=1, keep_dims=True) 88 | self.bots = tf.reduce_min(modes, axis=1, keep_dims=True) 89 | 90 | def get_parameters(self): 91 | return self.m + self.cov + [self.coef] 92 | 93 | def __call__(self, x): 94 | dist = [tf.contrib.distributions.MultivariateNormalDiag( 95 | self.m[i], self.cov[i], name='MultVarNorm_{}'.format(i)) 96 | for i in range(self.ncoef)] 97 | p = tf.concat([tf.reshape(dist[i].prob(x), [-1, 1]) 98 | for i in range(self.ncoef)], axis=1) 99 | 100 | fx = tf.matmul(p, self.coef) 101 | result = (fx / self.cst - self.bots) / (self.tops - self.bots) 102 | # import pdb; pdb.set_trace() 103 | if self.random: 104 | result = result + tf.random_normal(shape=[self.batch_size, 1], 105 | stddev=self.random, 106 | dtype=self.dtype, name='error') 107 | return result 108 | 109 | 110 | 111 | class Quadratic: 112 | """Quadratic problem: f(x) = ||Wx - y||.""" 113 | def __init__(self, batch_size=128, num_dims=3, ptype='convex', 114 | random=0.05, dtype=tf.float32): 115 | self.ptype = ptype 116 | self.w = tf.get_variable('w', shape=[batch_size, num_dims, num_dims], 117 | dtype=dtype, initializer=tf.random_normal_initializer(), 118 | trainable=False) 119 | 120 | self.a = tf.get_variable('y', shape=[batch_size, num_dims], 121 | dtype=dtype, initializer=tf.truncated_normal_initializer(mean=0.5, stddev=0.2), 122 | trainable=False) 123 | 124 | self.y = tf.squeeze(tf.matmul(self.w, tf.expand_dims(self.a, -1))) 125 | 126 | 127 | self.normalizer = tf.maximum( 128 | self._func(tf.zeros([batch_size, num_dims])), 129 | self._func(tf.ones([batch_size, num_dims]))) 130 | 131 | if random is not None: 132 | self.e = tf.random_normal(shape=[batch_size,], stddev=random, 133 | dtype=dtype, name='e') 134 | else: 135 | self.e = 0.0 136 | 137 | def get_parameters(self): 138 | return [self.w, self.a] 139 | 140 | def _func(self, var): 141 | product = tf.squeeze(tf.matmul(self.w, tf.expand_dims(var, -1))) 142 | norm = tf.reduce_sum((product - self.y) ** 2, 1) 143 | return norm 144 | 145 | def __call__(self, x): 146 | ''' 147 | x = tf.get_variable('x', shape=[batch_size, num_dims], 148 | dtype=dtype, initializer=tf.random_normal_initializer(stddev=stdev)) 149 | ''' 150 | res = (self._func(x) / self.normalizer + self.e) 151 | if self.ptype == 'concave': 152 | res = 1 - res 153 | return res 154 | 155 | 156 | class QuadraticEval: 157 | def __init__(self, num_dim=3, random=0.5, ptype='convex', 158 | dtype=np.float32, ifprint=False, record=False): 159 | self.ndim = num_dim 160 | self.dtype = dtype 161 | if random is not None: 162 | self.e = np.random.normal(scale=random) 163 | else: 164 | self.e = 0.0 165 | self.record = record 166 | self.refresh() 167 | self.normalizer = np.maximum( 168 | self._func(np.zeros([1, self.ndim], dtype=self.dtype)), 169 | self._func(np.ones([1, self.ndim], dtype=self.dtype))) 170 | self.ptype = ptype 171 | self.ifprint = ifprint 172 | 173 | def refresh(self): 174 | self.w = np.random.normal(size=(self.ndim, self.ndim)) 175 | self.a = np.random.uniform(low=0.01, high=0.99, size=(1, self.ndim)) 176 | self.y = np.dot(self.a, self.w) 177 | if self.record: 178 | self.history = {'x':[], 'y':[]} 179 | 180 | def _func(self, x): 181 | product = np.squeeze(np.dot(x, self.w)) 182 | norm = np.sum((product - self.y) ** 2) 183 | return norm 184 | 185 | def __call__(self, x): 186 | if self.ifprint: 187 | print('Input:') 188 | print(x) 189 | res = np.asscalar(self._func(x) / self.normalizer + self.e) 190 | if self.ptype == 'concave': 191 | res = 1 - res 192 | if self.ifprint: 193 | print('Output:') 194 | print(res) 195 | if self.record: 196 | self.history['x'].append(x) 197 | self.history['y'].append(res) 198 | return res 199 | 200 | class ConstraintQuadraticEval: 201 | def __init__(self, num_dim=3, random=0.5, ptype='convex', 202 | dtype=np.float32): 203 | self.ndim = num_dim 204 | self.dtype = dtype 205 | if random is not None: 206 | self.e = np.random.normal(scale=random) 207 | else: 208 | self.e = 0.0 209 | self.refresh() 210 | self.normalizer = np.maximum( 211 | self._func(np.zeros([1, self.ndim], dtype=self.dtype)), 212 | self._func(np.ones([1, self.ndim], dtype=self.dtype))) 213 | self.ptype = ptype 214 | 215 | def refresh(self): 216 | self.w = np.random.normal(size=(self.ndim, self.ndim)) 217 | self.a = np.maximum(np.minimum(np.random.normal(size=[self.ndim]), 0.8), 0.2) 218 | self.y = np.dot(self.a, self.w) 219 | 220 | def _func(self, x): 221 | product = np.squeeze(np.dot(x, self.w)) 222 | norm = np.sum((product - self.y) ** 2) 223 | return norm 224 | 225 | def _barrier(self, x): 226 | return - np.sum(np.log(x) + np.log(1-x), 1) / 1e10 227 | 228 | def __call__(self, x): 229 | print('Input:') 230 | print(x) 231 | res = np.asscalar(self._func(x) / self.normalizer + self.e + self._barrier(x)) 232 | if self.ptype == 'concave': 233 | res = 1 - res 234 | print('Output:') 235 | print(res) 236 | return res 237 | 238 | 239 | class RealReaction: 240 | def __init__(self, num_dim, param_range, param_names=['x1', 'x2', 'x3'], 241 | direction='max', logger=None): 242 | self.ndim = num_dim 243 | self.param_range = param_range 244 | self.param_names = param_names 245 | self.direction = direction 246 | 247 | def x_convert(self, x): 248 | real_x = np.zeros([self.ndim]) 249 | for i in range(self.ndim): 250 | a, b = self.param_range[i] 251 | real_x[i] = x[i] * (b - a) + a 252 | return real_x 253 | 254 | def y_convert(self, y): 255 | if self.direction == 'max': 256 | return 1 - y 257 | return y 258 | 259 | def __call__(self, x): 260 | print('Set Reaction Condition:') 261 | real_x = self.x_convert(np.squeeze(x)) 262 | for i in range(self.ndim): 263 | print('{0}: {1:.3f}'.format(self.param_names[i], real_x[i])) 264 | result = float(input('Input the reaction yield:')) 265 | return self.y_convert(result) 266 | 267 | 268 | -------------------------------------------------------------------------------- /realreaction.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import logging 4 | import matplotlib.pyplot as plt 5 | import json 6 | 7 | import rnn 8 | from reactions import QuadraticEval, ConstraintQuadraticEval, RealReaction 9 | from logger import get_handlers 10 | from collections import namedtuple 11 | 12 | logging.basicConfig(level=logging.INFO, handlers=get_handlers()) 13 | logger = logging.getLogger() 14 | 15 | 16 | class StepOptimizer: 17 | def __init__(self, cell, func, ndim, nsteps, ckpt_path, logger, constraints): 18 | self.logger = logger 19 | self.cell = cell 20 | self.func = func 21 | self.ndim = ndim 22 | self.nsteps = nsteps 23 | self.ckpt_path = ckpt_path 24 | self.constraints = constraints 25 | self.init_state = self.cell.get_initial_state(1, tf.float32) 26 | self.results = self.build_graph() 27 | 28 | self.saver = tf.train.Saver(tf.global_variables()) 29 | 30 | def get_state_shapes(self): 31 | return [(s[0].get_shape().as_list(), s[1].get_shape().as_list()) 32 | for s in self.init_state] 33 | 34 | def step(self, sess, x, y, state): 35 | feed_dict = {'input_x:0':x, 'input_y:0':y} 36 | for i in range(len(self.init_state)): 37 | feed_dict['state_l{0}_c:0'.format(i)] = state[i][0] 38 | feed_dict['state_l{0}_h:0'.format(i)] = state[i][1] 39 | new_x, new_state = sess.run(self.results, feed_dict=feed_dict) 40 | return new_x, new_state 41 | 42 | def build_graph(self): 43 | x = tf.placeholder(tf.float32, shape=[1, self.ndim], name='input_x') 44 | y = tf.placeholder(tf.float32, shape=[1, 1], name='input_y') 45 | state = [] 46 | for i in range(len(self.init_state)): 47 | state.append((tf.placeholder( 48 | tf.float32, shape=self.init_state[i][0].get_shape(), 49 | name='state_l{0}_c'.format(i)), 50 | tf.placeholder( 51 | tf.float32, shape=self.init_state[i][1].get_shape(), 52 | name='state_l{0}_h'.format(i)))) 53 | 54 | with tf.name_scope('opt_cell'): 55 | new_x, new_state = self.cell(x, y, state) 56 | if self.constraints: 57 | new_x = tf.clip_by_value(new_x, 0.01, 0.99) 58 | return new_x, new_state 59 | 60 | def load(self, sess, ckpt_path): 61 | ckpt = tf.train.get_checkpoint_state(ckpt_path) 62 | if ckpt and ckpt.model_checkpoint_path: 63 | logger.info('Reading model parameters from {}.'.format( 64 | ckpt.model_checkpoint_path)) 65 | self.saver.restore(sess, ckpt.model_checkpoint_path) 66 | else: 67 | raise FileNotFoundError('No checkpoint available') 68 | 69 | def get_init(self): 70 | x = np.random.normal(loc=0.5, scale=0.2, size=(1, 3)) 71 | x = np.maximum(np.minimum(x, 0.9), 0.1) 72 | y = np.array(self.func(x)).reshape(1, 1) 73 | init_state = [(np.zeros(s[0]), np.zeros(s[1])) 74 | for s in self.get_state_shapes()] 75 | return x, y, init_state 76 | 77 | def run(self): 78 | with tf.Session() as sess: 79 | self.load(sess, self.ckpt_path) 80 | x, y, state = self.get_init() 81 | x_array = np.zeros((self.nsteps + 1, self.ndim)) 82 | y_array = np.zeros((self.nsteps + 1, 1)) 83 | x_array[0, :] = x 84 | y_array[0] = y 85 | for i in range(self.nsteps): 86 | x, state = self.step(sess, x, y, state) 87 | y = np.array(self.func(x)).reshape(1, 1) 88 | x_array[i+1, :] = x 89 | y_array[i+1] = y 90 | 91 | return x_array, y_array 92 | 93 | def main(): 94 | config_file = open('./config.json') 95 | config = json.load(config_file, 96 | object_hook=lambda d:namedtuple('x', d.keys())(*d.values())) 97 | 98 | param_names = ['voltage', 'flow_rate', 'pressure'] 99 | param_range = [(0.0, 5.0), (1.0, 12.0), (10, 100)] 100 | func = RealReaction(num_dim = 3, param_range=param_range, param_names=param_names, 101 | direction='max', logger=None) 102 | 103 | cell = rnn.StochasticRNNCell(cell=rnn.LSTM, 104 | kwargs={'hidden_size':config.hidden_size}, 105 | nlayers=config.num_layers, 106 | reuse=config.reuse) 107 | optimizer = StepOptimizer(cell=cell, func=func, ndim=config.num_params, 108 | nsteps=config.num_steps, 109 | ckpt_path=config.save_path, logger=logger, 110 | constraints=config.constraints) 111 | x_array, y_array = optimizer.run() 112 | 113 | 114 | # plt.figure(1) 115 | # plt.plot(y_array) 116 | # plt.show() 117 | 118 | if __name__ == '__main__': 119 | main() 120 | -------------------------------------------------------------------------------- /rnn.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import batch_norm 3 | import util 4 | import pdb 5 | 6 | from tensorflow.contrib.rnn import RNNCell 7 | from tensorflow.python.ops import array_ops 8 | from tensorflow.contrib.rnn import LSTMCell, LSTMStateTuple 9 | from tensorflow.contrib.rnn import MultiRNNCell 10 | from tensorflow.python.ops.math_ops import tanh 11 | 12 | 13 | class MultiInputLSTM(LSTMCell): 14 | def __init__(self, nlayers, num_units, input_size=None, 15 | use_peepholes=False, cell_clip=None, initializer=None, 16 | num_proj=None, proj_clip=None, num_unit_shards=1, 17 | num_proj_shards=1, forget_bias=1.0, state_is_tuple=True, 18 | activation=tanh): 19 | 20 | super(MultiInputLSTM, self).__init__(num_units, input_size=None, 21 | use_peepholes=False,cell_clip=None, initializer=None, num_proj=None, 22 | proj_clip=None, num_unit_shards=1, num_proj_shards=1, 23 | forget_bias=1.0,state_is_tuple=True, activation=tanh) 24 | 25 | self.cell = super(MultiInputLSTM, self).__call__ 26 | 27 | if nlayers > 1: 28 | self.cell = MultiRNNCell([self.cell] * nlayers) 29 | self.nlayers = nlayers 30 | 31 | def __call__(self, x, y, state, scope=None): 32 | x_dim = x.get_shape()[1] 33 | inputs = tf.concat([x, tf.reshape(y, (-1, 1))], 34 | axis=1, name='inputs') 35 | output, nstate = self.cell(inputs, state) 36 | with tf.variable_scope('proj'): 37 | w = tf.get_variable('proj_weight', [self._num_units, x_dim]) 38 | x = tf.matmul(output, w) 39 | return x, nstate 40 | 41 | def get_initial_state(self, batch_size, dtype): 42 | if self.nlayers == 1: 43 | return super(MultiInputLSTM, self).zero_state(batch_size, dtype) 44 | else: 45 | return tuple( 46 | [super(MultiInputLSTM, self).zero_state(batch_size, dtype)] * self.nlayers) 47 | 48 | 49 | class MultiInputRNNCell(RNNCell): 50 | def __init__(self, cell, kwargs, nlayers=1, reuse=False): 51 | self.cell = cell(**kwargs, name="lstm") 52 | self.nlayers = nlayers 53 | self.rnncell = self.cell 54 | if nlayers > 1: 55 | if reuse: 56 | self.rnncell = MultiRNNCell([self.cell] * nlayers) 57 | else: 58 | self.rnncell = MultiRNNCell([cell(**kwargs, name='lstm_{}'.format(i)) 59 | for i in range(nlayers)]) 60 | 61 | def __call__(self, x, y, state, scope=None): 62 | with tf.variable_scope(scope or 'multi_input_rnn'): 63 | x_dim = int(x.get_shape()[1]) 64 | y = tf.tile(tf.reshape(y, [-1, 1]), [1, x_dim]) 65 | inputs = tf.concat([x, y], axis=1, name='inputs') 66 | output, nstate = self.rnncell(inputs, state) 67 | with tf.variable_scope('proj'): 68 | w = tf.get_variable('proj_weight', 69 | [self.cell.output_size.as_list()[0], x_dim]) 70 | b = tf.get_variable('proj_bias', [x_dim]) 71 | x = tf.matmul(output, w) + b 72 | return x, nstate 73 | 74 | def get_initial_state(self, batch_size, dtype=tf.float32): 75 | try: 76 | state = self.cell.get_initial_state(batch_size) 77 | except: 78 | state = self.cell.zero_state(batch_size, dtype) 79 | if self.nlayers == 1: 80 | return state 81 | else: 82 | return tuple([state] * self.nlayers) 83 | 84 | 85 | class StochasticRNNCell(RNNCell): 86 | def __init__(self, cell, kwargs, nlayers=1, reuse=False): 87 | self.cell = cell(**kwargs, name="lstm") 88 | self.nlayers = nlayers 89 | self.rnncell = self.cell 90 | if nlayers > 1: 91 | if reuse: 92 | self.rnncell = MultiRNNCell([self.cell] * nlayers) 93 | else: 94 | self.rnncell = MultiRNNCell([cell(**kwargs, name='lstm_{}'.format(i)) 95 | for i in range(nlayers)]) 96 | 97 | def __call__(self, x, y, state, scope=None): 98 | hidden_size = self.cell.output_size.as_list()[0] 99 | batch_size = x.get_shape().as_list()[0] 100 | with tf.variable_scope(scope or 'multi_input_rnn'): 101 | x_dim = int(x.get_shape()[1]) 102 | y = tf.tile(tf.reshape(y, [-1, 1]), [1, x_dim]) 103 | inputs = tf.concat([x, y], axis=1, name='inputs') 104 | output, nstate = self.rnncell(inputs, state) 105 | tot_dim = x_dim * (x_dim + 1) 106 | with tf.variable_scope('proj'): 107 | w = tf.get_variable('proj_weight', [hidden_size, tot_dim]) 108 | b = tf.get_variable('proj_bias', [tot_dim]) 109 | out = tf.matmul(output, w) + b 110 | mean, var = tf.split(out, [x_dim, x_dim ** 2], axis=1) 111 | var = tf.reshape(var, [batch_size, x_dim, x_dim]) 112 | dist = tf.contrib.distributions.MultivariateNormalTriL( 113 | mean, var, name='x_dist') 114 | x = dist.sample() 115 | 116 | return x, nstate 117 | 118 | def get_initial_state(self, batch_size, dtype=tf.float32): 119 | try: 120 | state = self.cell.get_initial_state(batch_size) 121 | except: 122 | state = self.cell.zero_state(batch_size, dtype) 123 | if self.nlayers == 1: 124 | return state 125 | else: 126 | return tuple([state] * self.nlayers) 127 | 128 | 129 | 130 | class LSTM(RNNCell): 131 | # Keys that may be provided for parameter initializers. 132 | W_GATES = "w_gates" # weight for gates 133 | B_GATES = "b_gates" # bias of gates 134 | W_F_DIAG = "w_f_diag" # weight for prev_cell -> forget gate peephole 135 | W_I_DIAG = "w_i_diag" # weight for prev_cell -> input gate peephole 136 | W_O_DIAG = "w_o_diag" # weight for prev_cell -> output gate peephole 137 | GAMMA_H = "gamma_h" # batch norm scaling for previous_hidden -> gates 138 | GAMMA_X = "gamma_x" # batch norm scaling for input -> gates 139 | GAMMA_C = "gamma_c" # batch norm scaling for cell -> output 140 | BETA_C = "beta_c" # (batch norm) bias for cell -> output 141 | POSSIBLE_KEYS = {W_GATES, B_GATES, W_F_DIAG, W_I_DIAG, W_O_DIAG, GAMMA_H, 142 | GAMMA_X, GAMMA_C, BETA_C} 143 | 144 | def __init__(self, 145 | hidden_size, 146 | forget_bias=1.0, 147 | initializers=None, 148 | use_peepholes=False, 149 | use_batch_norm_h=False, 150 | use_batch_norm_x=False, 151 | use_batch_norm_c=False, 152 | max_unique_stats=1, 153 | name="lstm"): 154 | super(LSTM, self).__init__() 155 | self.name_ = name 156 | self._template = tf.make_template(self.name_, self._build, 157 | create_scope_now_=True) 158 | self._hidden_size = hidden_size 159 | self._forget_bias = forget_bias 160 | self._use_peepholes = use_peepholes 161 | self._max_unique_stats = max_unique_stats 162 | self._use_batch_norm_h = use_batch_norm_h 163 | self._use_batch_norm_x = use_batch_norm_x 164 | self._use_batch_norm_c = use_batch_norm_c 165 | self.possible_keys = self.get_possible_initializer_keys(use_peepholes=use_peepholes, use_batch_norm_h=use_batch_norm_h, 166 | use_batch_norm_x=use_batch_norm_x, use_batch_norm_c=use_batch_norm_c) 167 | self._initializers = util.check_initializers(initializers, 168 | self.possible_keys) 169 | if max_unique_stats < 1: 170 | raise ValueError("max_unique_stats must be >= 1") 171 | if max_unique_stats != 1 and not ( 172 | use_batch_norm_h or use_batch_norm_x or use_batch_norm_c): 173 | raise ValueError("max_unique_stats specified but batch norm disabled") 174 | 175 | if use_batch_norm_h: 176 | self._batch_norm_h = LSTM.IndexedStatsBatchNorm(max_unique_stats, 177 | "batch_norm_h") 178 | if use_batch_norm_x: 179 | self._batch_norm_x = LSTM.IndexedStatsBatchNorm(max_unique_stats, 180 | "batch_norm_x") 181 | if use_batch_norm_c: 182 | self._batch_norm_c = LSTM.IndexedStatsBatchNorm(max_unique_stats, 183 | "batch_norm_c") 184 | 185 | def with_batch_norm_control(self, is_training=True, test_local_stats=True): 186 | return LSTM.CellWithExtraInput(self, 187 | is_training=is_training, 188 | test_local_stats=test_local_stats) 189 | 190 | @classmethod 191 | def get_possible_initializer_keys( 192 | cls, use_peepholes=False, use_batch_norm_h=False, use_batch_norm_x=False, 193 | use_batch_norm_c=False): 194 | possible_keys = cls.POSSIBLE_KEYS.copy() 195 | if not use_peepholes: 196 | possible_keys.difference_update( 197 | {cls.W_F_DIAG, cls.W_I_DIAG, cls.W_O_DIAG}) 198 | if not use_batch_norm_h: 199 | possible_keys.remove(cls.GAMMA_H) 200 | if not use_batch_norm_x: 201 | possible_keys.remove(cls.GAMMA_X) 202 | if not use_batch_norm_c: 203 | possible_keys.difference_update({cls.GAMMA_C, cls.BETA_C}) 204 | return possible_keys 205 | 206 | def _build(self, inputs, prev_state, is_training=True, test_local_stats=True): 207 | if self._max_unique_stats == 1: 208 | prev_hidden, prev_cell = prev_state 209 | time_step = None 210 | else: 211 | prev_hidden, prev_cell, time_step = prev_state 212 | 213 | self._create_gate_variables(inputs.get_shape(), inputs.dtype) 214 | self._create_batch_norm_variables(inputs.dtype) 215 | 216 | if self._use_batch_norm_h or self._use_batch_norm_x: 217 | gates_h = tf.matmul(prev_hidden, self._w_h) 218 | gates_x = tf.matmul(inputs, self._w_x) 219 | if self._use_batch_norm_h: 220 | gates_h = self._gamma_h * self._batch_norm_h(gates_h, 221 | time_step, 222 | is_training, 223 | test_local_stats) 224 | if self._use_batch_norm_x: 225 | gates_x = self._gamma_x * self._batch_norm_x(gates_x, 226 | time_step, 227 | is_training, 228 | test_local_stats) 229 | gates = gates_h + gates_x + self._b 230 | else: 231 | # Parameters of gates are concatenated into one multiply for efficiency. 232 | inputs_and_hidden = tf.concat([inputs, prev_hidden], axis=1) 233 | gates = tf.matmul(inputs_and_hidden, self._w_xh) + self._b 234 | 235 | # i = input_gate, j = new_input, f = forget_gate, o = output_gate 236 | i, j, f, o = array_ops.split(gates, 4, axis=1) 237 | 238 | if self._use_peepholes: # diagonal connections 239 | self._create_peephole_variables(inputs.dtype) 240 | f += self._w_f_diag * prev_cell 241 | i += self._w_i_diag * prev_cell 242 | 243 | forget_mask = tf.sigmoid(f + self._forget_bias) 244 | new_cell = forget_mask * prev_cell + tf.sigmoid(i) * tf.tanh(j) 245 | cell_output = new_cell 246 | if self._use_batch_norm_c: 247 | cell_output = (self._beta_c 248 | + self._gamma_c * self._batch_norm_c(cell_output, 249 | time_step, 250 | is_training, 251 | test_local_stats)) 252 | if self._use_peepholes: 253 | cell_output += self._w_o_diag * cell_output 254 | new_hidden = tf.tanh(cell_output) * tf.sigmoid(o) 255 | 256 | if self._max_unique_stats == 1: 257 | return new_hidden, (new_hidden, new_cell) 258 | else: 259 | return new_hidden, (new_hidden, new_cell, time_step + 1) 260 | 261 | def __call__(self, input, prev_state, 262 | is_training=True, test_local_stats=True): 263 | return self._template(input, prev_state, is_training, test_local_stats) 264 | 265 | def _create_batch_norm_variables(self, dtype): 266 | """Initialize the variables used for the `BatchNorm`s (if any).""" 267 | gamma_initializer = tf.constant_initializer(0.1) 268 | 269 | if self._use_batch_norm_h: 270 | self._gamma_h = tf.get_variable( 271 | LSTM.GAMMA_H, 272 | shape=[4 * self._hidden_size], 273 | dtype=dtype, 274 | initializer=(self._initializers.get(LSTM.GAMMA_H, gamma_initializer))) 275 | if self._use_batch_norm_x: 276 | self._gamma_x = tf.get_variable( 277 | LSTM.GAMMA_X, 278 | shape=[4 * self._hidden_size], 279 | dtype=dtype, 280 | initializer=(self._initializers.get(LSTM.GAMMA_X, gamma_initializer))) 281 | if self._use_batch_norm_c: 282 | self._gamma_c = tf.get_variable( 283 | LSTM.GAMMA_C, 284 | shape=[self._hidden_size], 285 | dtype=dtype, 286 | initializer=( 287 | self._initializers.get(LSTM.GAMMA_C, gamma_initializer))) 288 | self._beta_c = tf.get_variable( 289 | LSTM.BETA_C, 290 | shape=[self._hidden_size], 291 | dtype=dtype, 292 | initializer=self._initializers.get(LSTM.BETA_C)) 293 | 294 | def _create_gate_variables(self, input_shape, dtype): 295 | """Initialize the variables used for the gates.""" 296 | if len(input_shape) != 2: 297 | raise ValueError( 298 | "Rank of shape must be {} not: {}".format(2, len(input_shape))) 299 | input_size = input_shape.dims[1].value 300 | 301 | b_shape = [4 * self._hidden_size] 302 | 303 | equiv_input_size = self._hidden_size + input_size 304 | initializer = util.create_linear_initializer(equiv_input_size) 305 | 306 | if self._use_batch_norm_h or self._use_batch_norm_x: 307 | self._w_h = tf.get_variable( 308 | LSTM.W_GATES + "_H", 309 | shape=[self._hidden_size, 4 * self._hidden_size], 310 | dtype=dtype, 311 | initializer=self._initializers.get(LSTM.W_GATES, initializer)) 312 | self._w_x = tf.get_variable( 313 | LSTM.W_GATES + "_X", 314 | shape=[input_size, 4 * self._hidden_size], 315 | dtype=dtype, 316 | initializer=self._initializers.get(LSTM.W_GATES, initializer)) 317 | else: 318 | self._w_xh = tf.get_variable( 319 | LSTM.W_GATES, 320 | shape=[self._hidden_size + input_size, 4 * self._hidden_size], 321 | dtype=dtype, 322 | initializer=self._initializers.get(LSTM.W_GATES, initializer)) 323 | self._b = tf.get_variable( 324 | LSTM.B_GATES, 325 | shape=b_shape, 326 | dtype=dtype, 327 | initializer=self._initializers.get(LSTM.B_GATES, initializer)) 328 | 329 | def _create_peephole_variables(self, dtype): 330 | """Initialize the variables used for the peephole connections.""" 331 | self._w_f_diag = tf.get_variable( 332 | LSTM.W_F_DIAG, 333 | shape=[self._hidden_size], 334 | dtype=dtype, 335 | initializer=self._initializers.get(LSTM.W_F_DIAG)) 336 | self._w_i_diag = tf.get_variable( 337 | LSTM.W_I_DIAG, 338 | shape=[self._hidden_size], 339 | dtype=dtype, 340 | initializer=self._initializers.get(LSTM.W_I_DIAG)) 341 | self._w_o_diag = tf.get_variable( 342 | LSTM.W_O_DIAG, 343 | shape=[self._hidden_size], 344 | dtype=dtype, 345 | initializer=self._initializers.get(LSTM.W_O_DIAG)) 346 | 347 | def get_initial_state(self, batch_size, dtype=tf.float32, trainable=False, 348 | trainable_initializers=None): 349 | 350 | if self._max_unique_stats == 1: 351 | return super(LSTM, self).initial_state( 352 | batch_size, dtype, trainable, trainable_initializers) 353 | else: 354 | if not trainable: 355 | state = super(rnn_core.RNNCore, self).zero_state(batch_size, dtype) 356 | else: 357 | # We have to manually create the state ourselves so we don't create a 358 | # variable that never gets used for the third entry. 359 | state = util.trainable_initial_state( 360 | batch_size, 361 | (tf.TensorShape([self._hidden_size]), 362 | tf.TensorShape([self._hidden_size])), 363 | dtype, 364 | trainable_initializers) 365 | return (state[0], state[1], tf.constant(0, dtype=tf.int32)) 366 | 367 | @property 368 | def state_size(self): 369 | """Tuple of `tf.TensorShape`s indicating the size of state tensors.""" 370 | if self._max_unique_stats == 1: 371 | return (tf.TensorShape([self._hidden_size]), 372 | tf.TensorShape([self._hidden_size])) 373 | else: 374 | return (tf.TensorShape([self._hidden_size]), 375 | tf.TensorShape([self._hidden_size]), 376 | tf.TensorShape(1)) 377 | 378 | @property 379 | def output_size(self): 380 | """`tf.TensorShape` indicating the size of the core output.""" 381 | return tf.TensorShape([self._hidden_size]) 382 | 383 | @property 384 | def use_peepholes(self): 385 | """Boolean indicating whether peephole connections are used.""" 386 | return self._use_peepholes 387 | 388 | @property 389 | def use_batch_norm_h(self): 390 | """Boolean indicating whether batch norm for hidden -> gates is enabled.""" 391 | return self._use_batch_norm_h 392 | 393 | @property 394 | def use_batch_norm_x(self): 395 | """Boolean indicating whether batch norm for input -> gates is enabled.""" 396 | return self._use_batch_norm_x 397 | 398 | @property 399 | def use_batch_norm_c(self): 400 | """Boolean indicating whether batch norm for cell -> output is enabled.""" 401 | return self._use_batch_norm_c 402 | 403 | class IndexedStatsBatchNorm(object): 404 | def __init__(self, max_unique_stats, name=None): 405 | # super(LSTM.IndexedStatsBatchNorm, self).__init__() 406 | self._max_unique_stats = max_unique_stats 407 | 408 | def _build(self, inputs, index, is_training, test_local_stats): 409 | def create_batch_norm(): 410 | return batch_norm.BatchNorm(offset=False, scale=False)( 411 | inputs, is_training, test_local_stats) 412 | 413 | if self._max_unique_stats > 1: 414 | pred_fn_pairs = [(tf.equal(i, index), create_batch_norm) 415 | for i in range(self._max_unique_stats - 1)] 416 | out = tf.case(pred_fn_pairs, create_batch_norm) 417 | out.set_shape(inputs.get_shape()) # needed for tf.case shape inference 418 | return out 419 | else: 420 | return create_batch_norm() 421 | 422 | class CellWithExtraInput(RNNCell): 423 | def __init__(self, cell, *args, **kwargs): 424 | self._cell = cell 425 | self._args = args 426 | self._kwargs = kwargs 427 | 428 | def __call__(self, inputs, state): 429 | return self._cell(inputs, state, *self._args, **self._kwargs) 430 | 431 | @property 432 | def state_size(self): 433 | """Tuple indicating the size of nested state tensors.""" 434 | return self._cell.state_size 435 | 436 | @property 437 | def output_size(self): 438 | """`tf.TensorShape` indicating the size of the core output.""" 439 | return self._cell.output_size 440 | -------------------------------------------------------------------------------- /util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import math 3 | import reactions 4 | import tensorflow as tf 5 | 6 | import rnn 7 | from model import Optimizer 8 | from shutil import copyfile 9 | 10 | 11 | def run_epoch(sess, cost_op, ops, reset, num_unrolls): 12 | """Runs one optimization epoch.""" 13 | sess.run(reset) 14 | for _ in range(num_unrolls): 15 | results = sess.run([cost_op] + ops) 16 | return results[0], results[1:] 17 | 18 | def create_model(sess, config, logger): 19 | if not config.save_path == None: 20 | if not os.path.exists(config.save_path): 21 | os.mkdir(config.save_path) 22 | copyfile('config.json', os.path.join(config.save_path, 'config.json')) 23 | 24 | if config.opt_direction == 'max': 25 | problem_type = 'concave' 26 | else: 27 | problem_type = 'convex' 28 | 29 | if config.reaction_type == 'quad' and config.constraints == False: 30 | rxn_yeild = reactions.Quadratic( 31 | batch_size=config.batch_size, 32 | num_dims=config.num_params, 33 | ptype=problem_type, 34 | random=config.instrument_error) 35 | elif config.reaction_type == 'quad' and config.constraints == True: 36 | rxn_yeild = reactions.ConstraintQuadratic( 37 | batch_size=config.batch_size, 38 | num_dims=config.num_params, 39 | ptype=problem_type, 40 | random=config.instrument_error) 41 | elif config.reaction_type == 'gmm': 42 | rxn_yeild = reactions.GMM( 43 | batch_size=config.batch_size, 44 | num_dims=config.num_params, 45 | random=config.instrument_error, 46 | cov=config.norm_cov) 47 | if config.policy == 'srnn': 48 | cell = rnn.StochasticRNNCell(cell=rnn.LSTM, 49 | kwargs= 50 | {'hidden_size':config.hidden_size, 51 | 'use_batch_norm_h':config.batch_norm, 52 | 'use_batch_norm_x':config.batch_norm, 53 | 'use_batch_norm_c':config.batch_norm,}, 54 | nlayers=config.num_layers, 55 | reuse=config.reuse) 56 | if config.policy == 'rnn': 57 | cell = rnn.MultiInputRNNCell(cell=rnn.LSTM, 58 | kwargs= 59 | {'hidden_size':config.hidden_size, 60 | 'use_batch_norm_h':config.batch_norm, 61 | 'use_batch_norm_x':config.batch_norm, 62 | 'use_batch_norm_c':config.batch_norm,}, 63 | nlayers=config.num_layers, 64 | reuse=config.reuse) 65 | model = Optimizer(cell=cell, logger=logger, func=rxn_yeild, 66 | ndim=config.num_params, batch_size=config.batch_size, 67 | unroll_len=config.unroll_length, lr=config.learning_rate, 68 | loss_type=config.loss_type, optimizer=config.optimizer, 69 | trainable_init=config.trainable_init, 70 | direction=config.opt_direction, constraints=config.constraints, 71 | discount_factor=config.discount_factor) 72 | 73 | ckpt = tf.train.get_checkpoint_state(config.save_path) 74 | if ckpt and ckpt.model_checkpoint_path: 75 | logger.info('Reading model parameters from {}.'.format( 76 | ckpt.model_checkpoint_path)) 77 | model.saver.restore(sess, ckpt.model_checkpoint_path) 78 | else: 79 | logger.info('Creating Model with fresh parameters.') 80 | sess.run(tf.global_variables_initializer()) 81 | return model 82 | 83 | def load_model(sess, config, logger): 84 | assert(os.path.exists(config.save_path)) 85 | 86 | if config.opt_direction == 'max': 87 | problem_type = 'concave' 88 | else: 89 | problem_type = 'convex' 90 | 91 | if config.reaction_type == 'quad' and config.constraints == False: 92 | rxn_yeild = reactions.Quadratic( 93 | batch_size=config.batch_size, 94 | num_dims=config.num_params, 95 | ptype=problem_type, 96 | random=config.instrument_error) 97 | elif config.reaction_type == 'quad' and config.constraints == True: 98 | rxn_yeild = reactions.ConstraintQuadratic( 99 | batch_size=config.batch_size, 100 | num_dims=config.num_params, 101 | ptype=problem_type, 102 | random=config.instrument_error) 103 | elif config.reaction_type == 'gmm': 104 | rxn_yeild = reactions.GMM( 105 | batch_size=config.batch_size, 106 | num_dims=config.num_params, 107 | random=config.instrument_error, 108 | cov=config.norm_cov) 109 | 110 | if config.policy == 'srnn': 111 | cell = rnn.StochasticRNNCell(cell=rnn.LSTM, 112 | kwargs= 113 | {'hidden_size':config.hidden_size, 114 | 'use_batch_norm_h':config.batch_norm, 115 | 'use_batch_norm_x':config.batch_norm, 116 | 'use_batch_norm_c':config.batch_norm,}, 117 | nlayers=config.num_layers, 118 | reuse=config.reuse) 119 | if config.policy == 'rnn': 120 | cell = rnn.MultiInputRNNCell(cell=rnn.LSTM, 121 | kwargs= 122 | {'hidden_size':config.hidden_size, 123 | 'use_batch_norm_h':config.batch_norm, 124 | 'use_batch_norm_x':config.batch_norm, 125 | 'use_batch_norm_c':config.batch_norm,}, 126 | nlayers=config.num_layers, 127 | reuse=config.reuse) 128 | model = Optimizer(cell=cell, logger=logger, func=rxn_yeild, 129 | ndim=config.num_params, batch_size=config.batch_size, 130 | unroll_len=config.unroll_length, lr=config.learning_rate, 131 | loss_type=config.loss_type, optimizer=config.optimizer, 132 | trainable_init=config.trainable_init, 133 | direction=config.opt_direction, constraints=config.constraints, 134 | discount_factor=config.discount_factor) 135 | 136 | 137 | ckpt = tf.train.get_checkpoint_state(config.save_path) 138 | if ckpt and ckpt.model_checkpoint_path: 139 | logger.info('Reading model parameters from {}.'.format( 140 | ckpt.model_checkpoint_path)) 141 | model.saver.restore(sess, ckpt.model_checkpoint_path) 142 | 143 | return model 144 | 145 | 146 | def check_initializers(initializers, keys): 147 | if initializers is None: 148 | return {} 149 | keys = set(keys) 150 | 151 | if not issubclass(type(initializers), dict): 152 | raise TypeError("A dict of initializers was expected, but not " 153 | "given. You should double-check that you've nested the " 154 | "initializers for any sub-modules correctly.") 155 | 156 | if not set(initializers) <= keys: 157 | extra_keys = set(initializers) - keys 158 | raise KeyError( 159 | "Invalid initializer keys {}, initializers can only " 160 | "be provided for {}".format( 161 | ", ".join("'{}'".format(key) for key in extra_keys), 162 | ", ".join("'{}'".format(key) for key in keys))) 163 | 164 | def check_nested_callables(dictionary): 165 | for key, entry in dictionary.items(): 166 | if isinstance(entry, dict): 167 | check_nested_callables(entry) 168 | elif not callable(entry): 169 | raise TypeError( 170 | "Initializer for '{}' is not a callable function " 171 | "or dictionary".format(key)) 172 | check_nested_callables(initializers) 173 | return dict(initializers) 174 | 175 | def create_linear_initializer(input_size): 176 | """Returns a default initializer for weights or bias of a linear module.""" 177 | stddev = 1 / math.sqrt(input_size) 178 | return tf.truncated_normal_initializer(stddev=stddev) 179 | 180 | def trainable_initial_state(batch_size, state_size, dtype, initializers=None): 181 | flat_state_size = nest.flatten(state_size) 182 | 183 | if not initializers: 184 | flat_initializer = tuple(tf.zeros_initializer for _ in flat_state_size) 185 | else: 186 | nest.assert_same_structure(initializers, state_size) 187 | flat_initializer = nest.flatten(initializers) 188 | if not all([callable(init) for init in flat_initializer]): 189 | raise ValueError("Not all the passed initializers are callable objects.") 190 | 191 | # Produce names for the variables. In the case of a tuple or nested tuple, 192 | # this is just a sequence of numbers, but for a flat `namedtuple`, we use 193 | # the field names. NOTE: this could be extended to nested `namedtuple`s, 194 | # but for now that's extra complexity that's not used anywhere. 195 | try: 196 | names = ["init_{}".format(state_size._fields[i]) 197 | for i in range(len(flat_state_size))] 198 | except (AttributeError, IndexError): 199 | names = ["init_state_{}".format(i) for i in range(len(flat_state_size))] 200 | 201 | flat_initial_state = [] 202 | 203 | for name, size, init in zip(names, flat_state_size, flat_initializer): 204 | shape_with_batch_dim = [1] + tensor_shape.as_shape(size).as_list() 205 | initial_state_variable = tf.get_variable( 206 | name, shape=shape_with_batch_dim, dtype=dtype, initializer=init) 207 | 208 | initial_state_variable_dims = initial_state_variable.get_shape().ndims 209 | tile_dims = [batch_size] + [1] * (initial_state_variable_dims - 1) 210 | flat_initial_state.append( 211 | tf.tile(initial_state_variable, tile_dims, name=(name + "_tiled"))) 212 | 213 | return nest.pack_sequence_as(structure=state_size, 214 | flat_sequence=flat_initial_state) 215 | --------------------------------------------------------------------------------