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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_LSTMStep_pickle(): """Tests that LSTMStep can be pickled.""" n_feat = 10 tg = TensorGraph() y = Feature(shape=(None, 2 * n_feat)) state_zero = Feature(shape=(None, n_feat)) state_one = Feature(shape=(None, n_feat)) lstm = LSTMStep(n_feat, 2 * n_feat, in_layers=[y, state_zero, state_one]) tg.add_output(lstm) tg.set_loss(lstm) tg.build() tg.save()
def test_Weave_pickle():
tg = TensorGraph()
atom_feature = Feature(shape=(None, 75))
pair_feature = Feature(shape=(None, 14))
pair_split = Feature(shape=(None,), dtype=tf.int32)
atom_to_pair = Feature(shape=(None, 2), dtype=tf.int32)
weave = WeaveLayer(
in_layers=[atom_feature, pair_feature, pair_split, atom_to_pair])
tg.add_output(weave)
tg.set_loss(weave)
tg.build()
tg.save()
def test_AtomicDifferentialDense_pickle():
max_atoms = 23
atom_features = 100
tg = TensorGraph()
atom_feature = Feature(shape=(None, max_atoms, atom_features))
atom_numbers = Feature(shape=(None, max_atoms))
atomic_differential_dense = AtomicDifferentiatedDense(
max_atoms=23, out_channels=5, in_layers=[atom_feature, atom_numbers])
tg.add_output(atomic_differential_dense)
tg.set_loss(atomic_differential_dense)
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)
value = policy_layers['value']
rewards = Weights(shape=(None, ))
advantages = Weights(shape=(None, ))
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)
if 'action_prob' in policy_layers:
self.continuous = False
action_prob = policy_layers['action_prob']
actions = Label(shape=(None, self._env.n_actions))
loss = A3CLossDiscrete(
self.value_weight,
self.entropy_weight,
in_layers=[rewards, actions, action_prob, value, advantages])
graph.add_output(action_prob)
else:
self.continuous = True
action_mean = policy_layers['action_mean']
action_std = policy_layers['action_std']
actions = Label(shape=[None] + list(self._env.action_shape))
loss = A3CLossContinuous(self.value_weight,
self.entropy_weight,
in_layers=[
rewards, actions, action_mean,
action_std, value, advantages
])
graph.add_output(action_mean)
graph.add_output(action_std)
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 self.continuous:
return graph, features, rewards, actions, action_mean, action_std, value, advantages
else:
return graph, features, rewards, actions, action_prob, value, advantages
def test_IRVLayer_pickle():
n_tasks = 10
K = 10
V = Feature(shape=(None, 200))
irv_layer = IRVLayer(n_tasks, K, in_layers=[V])
irv_reg = IRVRegularize(irv_layer, in_layers=[irv_layer])
tg = TensorGraph()
tg.add_output(irv_layer)
tg.add_output(irv_reg)
tg.set_loss(irv_reg)
tg.build()
tg.save()
def test_DTNNStep_pickle():
tg = TensorGraph()
atom_features = Feature(shape=(None, 30))
distance = Feature(shape=(None, 100))
distance_membership_i = Feature(shape=(None,), dtype=tf.int32)
distance_membership_j = Feature(shape=(None,), dtype=tf.int32)
DTNN = DTNNStep(in_layers=[
atom_features, distance, distance_membership_i, distance_membership_j
])
tg.add_output(DTNN)
tg.set_loss(DTNN)
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
