def _build_graph(self, tf_graph, scope, model_dir): """Construct a TensorGraph containing the policy and loss calculations.""" state_shape = self._env.state_shape state_dtype = self._env.state_dtype if not self._state_is_list: state_shape = [state_shape] state_dtype = [state_dtype] features = [] for s, d in zip(state_shape, state_dtype): features.append(Feature(shape=[None] + list(s), dtype=tf.as_dtype(d))) policy_layers = self._policy.create_layers(features) action_prob = policy_layers['action_prob'] value = policy_layers['value'] search_prob = Label(shape=(None, self._env.n_actions)) search_value = Label(shape=(None,)) loss = MCTSLoss( self.value_weight, in_layers=[action_prob, value, search_prob, search_value]) graph = TensorGraph( batch_size=self.max_search_depth, use_queue=False, graph=tf_graph, model_dir=model_dir) for f in features: graph._add_layer(f) graph.add_output(action_prob) graph.add_output(value) graph.set_loss(loss) graph.set_optimizer(self._optimizer) with graph._get_tf("Graph").as_default(): with tf.variable_scope(scope): graph.build() if len(graph.rnn_initial_states) > 0: raise ValueError('MCTS does not support policies with recurrent layers') return graph, features, action_prob, value, search_prob, search_value
def test_Conv1D_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1, 1)) conv = Conv1D(2, 1, in_layers=feature) tg.add_output(conv) tg.set_loss(conv) tg.build() tg.save()
def test_LSTM_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 10, 10)) layer = LSTM(n_hidden=10, batch_size=tg.batch_size, in_layers=feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_DTNNExtract_pickle(): tg = TensorGraph() atom_features = Feature(shape=(None, 30)) Ext = DTNNExtract(0, in_layers=[atom_features]) tg.add_output(Ext) tg.set_loss(Ext) tg.build() tg.save()
def test_Squeeze_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Squeeze(in_layers=feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_Dense_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) dense = Dense(out_channels=1, in_layers=feature) tg.add_output(dense) tg.set_loss(dense) tg.build() tg.save()
def test_Exp_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Exp(feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_DTNNEmbedding_pickle(): tg = TensorGraph() atom_numbers = Feature(shape=(None, 23), dtype=tf.int32) Embedding = DTNNEmbedding(in_layers=[atom_numbers]) tg.add_output(Embedding) tg.set_loss(Embedding) tg.build() tg.save()
def test_CombineMeanStd_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = CombineMeanStd(in_layers=[feature, feature]) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_Cast_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Cast(in_layers=feature, dtype=tf.int32) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_Gather_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Gather(indices=[[0], [2], [3]], in_layers=feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_StopGradient_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) output = StopGradient(feature) tg.add_output(output) tg.set_loss(output) tg.build() tg.save()
def test_Transpose_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Transpose(perm=(1, 0), in_layers=feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_SigmoidCrossEntropy_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = SigmoidCrossEntropy(in_layers=[feature, feature]) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_Reshape_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Reshape(shape=(None, 2), in_layers=feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_Repeat_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Repeat(n_times=10, in_layers=feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_BatchNorm_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 10)) layer = BatchNorm(in_layers=feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_WeightedError_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 10)) layer = WeightedError(in_layers=[feature, feature]) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_Conv3DTranspose_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 10, 10, 10, 1)) layer = Conv3DTranspose(num_outputs=3, in_layers=feature) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_ReduceSquareDifference_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = ReduceSquareDifference(in_layers=[feature, feature]) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_ToFloat_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = ToFloat(in_layers=[feature]) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_Slice_pickle(): V = Feature(shape=(None, 10)) out = Slice(5, 1, in_layers=[V]) tg = TensorGraph() tg.add_output(out) tg.set_loss(out) tg.build() tg.save()
def test_hingeloss_pickle(): tg = TensorGraph() feature = Feature(shape=(1, None)) layer = HingeLoss(in_layers=[feature, feature]) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def testGraphCNNPoolLayer_pickle(): V = Feature(shape=(None, 200, 50)) A = Feature(shape=(None, 200, 1, 200)) gcnnpool = GraphEmbedPoolLayer(32, in_layers=[V, A]) tg = TensorGraph() tg.add_output(gcnnpool) tg.set_loss(gcnnpool) tg.build() tg.save()
def testGraphCNN_pickle(): V = Feature(shape=(None, 200, 50)) A = Feature(shape=(None, 200, 1, 200)) gcnn = GraphCNN(32, in_layers=[V, A]) tg = TensorGraph() tg.add_output(gcnn) tg.set_loss(gcnn) tg.build() tg.save()
def test_Constant_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Constant(np.array([15.0])) output = Add(in_layers=[feature, layer]) tg.add_output(output) tg.set_loss(output) tg.build() tg.save()
def test_DAGGather_pickle(): tg = TensorGraph() atom_features = Feature(shape=(None, 30)) membership = Feature(shape=(None,), dtype=tf.int32) Gather = DAGGather(in_layers=[atom_features, membership]) tg.add_output(Gather) tg.set_loss(Gather) tg.build() tg.save()
def test_SetGather_pickle(): tg = TensorGraph() atom_feature = Feature(shape=(None, 100)) atom_split = Feature(shape=(None,), dtype=tf.int32) Gather = SetGather(5, 16, in_layers=[atom_feature, atom_split]) tg.add_output(Gather) tg.set_loss(Gather) tg.build() tg.save()
def test_SparseSoftmaxCrossEntropy_pickle(): tg = TensorGraph() logits = Feature(shape=(tg.batch_size, 5)) labels = Feature(shape=(tg.batch_size,), dtype=tf.int32) layer = SparseSoftMaxCrossEntropy(in_layers=[labels, logits]) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def test_Variable_pickle(): tg = TensorGraph() feature = Feature(shape=(tg.batch_size, 1)) layer = Variable(np.array([15.0])) output = Multiply(in_layers=[feature, layer]) tg.add_output(output) tg.set_loss(output) tg.build() tg.save()
def test_DAGLayer_pickle(): tg = TensorGraph(use_queue=False) atom_features = Feature(shape=(None, 75)) parents = Feature(shape=(None, 50, 50), dtype=tf.int32) calculation_orders = Feature(shape=(None, 50), dtype=tf.int32) calculation_masks = Feature(shape=(None, 50), dtype=tf.bool) n_atoms = Feature(shape=(), dtype=tf.int32) DAG = DAGLayer(in_layers=[ atom_features, parents, calculation_orders, calculation_masks, n_atoms ]) tg.add_output(DAG) tg.set_loss(DAG) tg.build() tg.save()
def test_GraphPool_Pickle(): tg = TensorGraph() atom_features = Feature(shape=(None, 75)) degree_slice = Feature(shape=(None, 2), dtype=tf.int32) membership = Feature(shape=(None, ), dtype=tf.int32) deg_adjs = [] for i in range(0, 10 + 1): deg_adj = Feature(shape=(None, i + 1), dtype=tf.int32) deg_adjs.append(deg_adj) layer = GraphPool(in_layers=[atom_features, degree_slice, membership] + deg_adjs) tg.set_loss(layer) tg.build() tg.save()
def test_IterRefLSTM_pickle(): """Tests that IterRefLSTM can be pickled.""" n_feat = 10 max_depth = 5 n_test = 5 n_support = 5 tg = TensorGraph() test = Feature(shape=(None, n_feat)) support = Feature(shape=(None, n_feat)) lstm = IterRefLSTMEmbedding( n_test, n_support, n_feat, max_depth, in_layers=[test, support]) tg.add_output(lstm) tg.set_loss(lstm) tg.build() tg.save()
def test_AttnLSTM_pickle(): """Tests that AttnLSTM can be pickled.""" max_depth = 5 n_test = 5 n_support = 5 n_feat = 10 tg = TensorGraph(batch_size=n_test) test = Feature(shape=(None, n_feat)) support = Feature(shape=(None, n_feat)) out = AttnLSTMEmbedding( n_test, n_support, n_feat, max_depth, in_layers=[test, support]) tg.add_output(out) tg.set_loss(out) tg.build() tg.save()
def test_GraphGather_Pickle(): tg = TensorGraph() atom_features = Feature(shape=(None, 75)) degree_slice = Feature(shape=(None, 2), dtype=tf.int32) membership = Feature(shape=(None,), dtype=tf.int32) deg_adjs = [] for i in range(0, 10 + 1): deg_adj = Feature(shape=(None, i + 1), dtype=tf.int32) deg_adjs.append(deg_adj) layer = GraphGather( batch_size=tg.batch_size, activation_fn=tf.nn.tanh, in_layers=[atom_features, degree_slice, membership] + deg_adjs) tg.set_loss(layer) tg.build() tg.save()
def test_GraphConv_pickle(): tg = TensorGraph() atom_features = Feature(shape=(None, 75)) degree_slice = Feature(shape=(None, 2), dtype=tf.int32) membership = Feature(shape=(None,), dtype=tf.int32) deg_adjs = [] for i in range(0, 10 + 1): deg_adj = Feature(shape=(None, i + 1), dtype=tf.int32) deg_adjs.append(deg_adj) layer = GraphConv( 64, activation_fn=tf.nn.relu, in_layers=[atom_features, degree_slice, membership] + deg_adjs) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def testInteratomicL2Distances(): """ TODO(LESWING) what is ndim here? :return: """ tg = TensorGraph() n_atoms = tg.batch_size M_nbrs = 4 n_dim = 3 feature = Feature(shape=(tg.batch_size, 3)) neighbors = Feature(shape=(tg.batch_size, M_nbrs), dtype=tf.int32) layer = InteratomicL2Distances(N_atoms=n_atoms, M_nbrs=M_nbrs, ndim=n_dim, in_layers=[feature, neighbors]) tg.add_output(layer) tg.set_loss(layer) tg.build() tg.save()
def _build_graph(self, tf_graph, scope, model_dir): """Construct a TensorGraph containing the policy and loss calculations.""" state_shape = self._env.state_shape state_dtype = self._env.state_dtype if not self._state_is_list: state_shape = [state_shape] state_dtype = [state_dtype] features = [] for s, d in zip(state_shape, state_dtype): features.append(Feature(shape=[None] + list(s), dtype=tf.as_dtype(d))) policy_layers = self._policy.create_layers(features) action_prob = policy_layers['action_prob'] value = policy_layers['value'] rewards = Weights(shape=(None,)) advantages = Weights(shape=(None,)) old_action_prob = Weights(shape=(None,)) actions = Label(shape=(None, self._env.n_actions)) loss = PPOLoss( self.value_weight, self.entropy_weight, self.clipping_width, in_layers=[ rewards, actions, action_prob, value, advantages, old_action_prob ]) graph = TensorGraph( batch_size=self.max_rollout_length, use_queue=False, graph=tf_graph, model_dir=model_dir) for f in features: graph._add_layer(f) graph.add_output(action_prob) graph.add_output(value) graph.set_loss(loss) graph.set_optimizer(self._optimizer) with graph._get_tf("Graph").as_default(): with tf.variable_scope(scope): graph.build() assert len(loss.components) > 0 return graph, features, rewards, actions, action_prob, value, advantages, old_action_prob, loss.components
def _build_graph(self, tf_graph, scope, model_dir): """Construct a TensorGraph containing the policy and loss calculations.""" state_shape = self._env.state_shape state_dtype = self._env.state_dtype if not self._state_is_list: state_shape = [state_shape] state_dtype = [state_dtype] features = [] for s, d in zip(state_shape, state_dtype): features.append( Feature(shape=[None] + list(s), dtype=tf.as_dtype(d))) policy_layers = self._policy.create_layers(features) action_prob = policy_layers['action_prob'] value = policy_layers['value'] search_prob = Label(shape=(None, self._env.n_actions)) search_value = Label(shape=(None, )) loss = MCTSLoss( self.value_weight, in_layers=[action_prob, value, search_prob, search_value]) graph = TensorGraph(batch_size=self.max_search_depth, use_queue=False, graph=tf_graph, model_dir=model_dir) for f in features: graph._add_layer(f) graph.add_output(action_prob) graph.add_output(value) graph.set_loss(loss) graph.set_optimizer(self._optimizer) with graph._get_tf("Graph").as_default(): with tf.variable_scope(scope): graph.build() if len(graph.rnn_initial_states) > 0: raise ValueError( 'MCTS does not support policies with recurrent layers') return graph, features, action_prob, value, search_prob, search_value
def atomic_conv_model(frag1_num_atoms=70, frag2_num_atoms=634, complex_num_atoms=701, max_num_neighbors=12, batch_size=24, at=[ 6, 7., 8., 9., 11., 12., 15., 16., 17., 20., 25., 30., 35., 53., -1. ], radial=[[ 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0, 11.5, 12.0 ], [0.0, 4.0, 8.0], [0.4]], layer_sizes=[32, 32, 16], learning_rate=0.001): rp = [x for x in itertools.product(*radial)] frag1_X = Feature(shape=(batch_size, frag1_num_atoms, 3)) frag1_nbrs = Feature(shape=(batch_size, frag1_num_atoms, max_num_neighbors)) frag1_nbrs_z = Feature(shape=(batch_size, frag1_num_atoms, max_num_neighbors)) frag1_z = Feature(shape=(batch_size, frag1_num_atoms)) frag2_X = Feature(shape=(batch_size, frag2_num_atoms, 3)) frag2_nbrs = Feature(shape=(batch_size, frag2_num_atoms, max_num_neighbors)) frag2_nbrs_z = Feature(shape=(batch_size, frag2_num_atoms, max_num_neighbors)) frag2_z = Feature(shape=(batch_size, frag2_num_atoms)) complex_X = Feature(shape=(batch_size, complex_num_atoms, 3)) complex_nbrs = Feature(shape=(batch_size, complex_num_atoms, max_num_neighbors)) complex_nbrs_z = Feature(shape=(batch_size, complex_num_atoms, max_num_neighbors)) complex_z = Feature(shape=(batch_size, complex_num_atoms)) frag1_conv = AtomicConvolution( atom_types=at, radial_params=rp, boxsize=None, in_layers=[frag1_X, frag1_nbrs, frag1_nbrs_z]) frag2_conv = AtomicConvolution( atom_types=at, radial_params=rp, boxsize=None, in_layers=[frag2_X, frag2_nbrs, frag2_nbrs_z]) complex_conv = AtomicConvolution( atom_types=at, radial_params=rp, boxsize=None, in_layers=[complex_X, complex_nbrs, complex_nbrs_z]) score = AtomicConvScore(at, layer_sizes, in_layers=[ frag1_conv, frag2_conv, complex_conv, frag1_z, frag2_z, complex_z ]) label = Label(shape=(None, 1)) loss = ReduceMean(in_layers=L2Loss(in_layers=[score, label])) def feed_dict_generator(dataset, batch_size, epochs=1, pad_batches=True): def replace_atom_types(z): def place_holder(i): if i in at: return i return -1 return np.array([place_holder(x) for x in z]) for epoch in range(epochs): for ind, (F_b, y_b, w_b, ids_b) in enumerate( dataset.iterbatches(batch_size, deterministic=True, pad_batches=pad_batches)): N = complex_num_atoms N_1 = frag1_num_atoms N_2 = frag2_num_atoms M = max_num_neighbors orig_dict = {} batch_size = F_b.shape[0] num_features = F_b[0][0].shape[1] frag1_X_b = np.zeros((batch_size, N_1, num_features)) for i in range(batch_size): frag1_X_b[i] = F_b[i][0] orig_dict[frag1_X] = frag1_X_b frag2_X_b = np.zeros((batch_size, N_2, num_features)) for i in range(batch_size): frag2_X_b[i] = F_b[i][3] orig_dict[frag2_X] = frag2_X_b complex_X_b = np.zeros((batch_size, N, num_features)) for i in range(batch_size): complex_X_b[i] = F_b[i][6] orig_dict[complex_X] = complex_X_b frag1_Nbrs = np.zeros((batch_size, N_1, M)) frag1_Z_b = np.zeros((batch_size, N_1)) for i in range(batch_size): z = replace_atom_types(F_b[i][2]) frag1_Z_b[i] = z frag1_Nbrs_Z = np.zeros((batch_size, N_1, M)) for atom in range(N_1): for i in range(batch_size): atom_nbrs = F_b[i][1].get(atom, "") frag1_Nbrs[i, atom, :len(atom_nbrs)] = np.array(atom_nbrs) for j, atom_j in enumerate(atom_nbrs): frag1_Nbrs_Z[i, atom, j] = frag1_Z_b[i, atom_j] orig_dict[frag1_nbrs] = frag1_Nbrs orig_dict[frag1_nbrs_z] = frag1_Nbrs_Z orig_dict[frag1_z] = frag1_Z_b frag2_Nbrs = np.zeros((batch_size, N_2, M)) frag2_Z_b = np.zeros((batch_size, N_2)) for i in range(batch_size): z = replace_atom_types(F_b[i][5]) frag2_Z_b[i] = z frag2_Nbrs_Z = np.zeros((batch_size, N_2, M)) for atom in range(N_2): for i in range(batch_size): atom_nbrs = F_b[i][4].get(atom, "") frag2_Nbrs[i, atom, :len(atom_nbrs)] = np.array(atom_nbrs) for j, atom_j in enumerate(atom_nbrs): frag2_Nbrs_Z[i, atom, j] = frag2_Z_b[i, atom_j] orig_dict[frag2_nbrs] = frag2_Nbrs orig_dict[frag2_nbrs_z] = frag2_Nbrs_Z orig_dict[frag2_z] = frag2_Z_b complex_Nbrs = np.zeros((batch_size, N, M)) complex_Z_b = np.zeros((batch_size, N)) for i in range(batch_size): z = replace_atom_types(F_b[i][8]) complex_Z_b[i] = z complex_Nbrs_Z = np.zeros((batch_size, N, M)) for atom in range(N): for i in range(batch_size): atom_nbrs = F_b[i][7].get(atom, "") complex_Nbrs[i, atom, :len(atom_nbrs)] = np.array( atom_nbrs) for j, atom_j in enumerate(atom_nbrs): complex_Nbrs_Z[i, atom, j] = complex_Z_b[i, atom_j] orig_dict[complex_nbrs] = complex_Nbrs orig_dict[complex_nbrs_z] = complex_Nbrs_Z orig_dict[complex_z] = complex_Z_b orig_dict[label] = np.reshape(y_b, newshape=(batch_size, 1)) yield orig_dict tg = TensorGraph(batch_size=batch_size, mode=str("regression"), model_dir=str("/tmp/atom_conv"), learning_rate=learning_rate) tg.add_output(score) tg.set_loss(loss) return tg, feed_dict_generator, label