Can pickle AttnLSTM (ca077d31) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/layers.py

+37 −8

Original line number	Diff line number	Diff line
		@@ -91,6 +91,8 @@ class Layer(object):
		for input in in_layers:
		if isinstance(input, tf.Tensor):
		tensors.append(input)
		elif isinstance(input, tf.Variable):
		tensors.append(input)
		elif isinstance(input, Layer):
		tensors.append(input.out_tensor)
		else:
		@@ -1445,8 +1447,8 @@ class LSTMStep(Layer):
		inner_activation = activations.get(self.inner_activation)

		self.build()
		if in_layers is None:
		in_layers = self.in_layers
		#if in_layers is None:
		# in_layers = self.in_layers
		inputs = self._get_input_tensors(in_layers)
		x, h_tm1, c_tm1 = inputs

		@@ -1465,7 +1467,7 @@ class LSTMStep(Layer):

		h = o * activation(c)

		# TODO(rbharath): Is this correct?
		# TODO(rbharath): This is wrong!!
		self.out_tensor = h
		return h, [h, c]

		@@ -1520,7 +1522,7 @@ class AttnLSTMEmbedding(Layer):
		self.n_support = n_support
		self.n_feat = n_feat

		def create_tensor(self, in_layers=None, set_tensors=True):
		def create_tensor(self, in_layers=None, **kwargs):
		"""Execute this layer on input tensors.

		Parameters
		@@ -1546,10 +1548,10 @@ class AttnLSTMEmbedding(Layer):
		## Initializes trainable weights.
		n_feat = self.n_feat

		self.lstm = LSTMStep(n_feat, 2 * n_feat)
		lstm = LSTMStep(n_feat, 2 * n_feat)
		self.q_init = model_ops.zeros([self.n_test, n_feat])
		self.r_init = model_ops.zeros([self.n_test, n_feat])
		self.states_init = self.lstm.get_initial_states([self.n_test, n_feat])
		self.states_init = lstm.get_initial_states([self.n_test, n_feat])

		self.trainable_weights = [self.q_init, self.r_init]

		@@ -1568,10 +1570,25 @@ class AttnLSTMEmbedding(Layer):

		# Generate new aattention states
		y = model_ops.concatenate([q, r], axis=1)
		q, states = self.lstm(y, *states) #+ self.lstm.get_constants(x)
		q, states = lstm(y, *states)

		# TODO(rbharath): This is wrong!!
		self.out_tensor = xp
		return [x + q, xp]

		def none_tensors(self):
		q_init, r_init, states_init = self.q_init, self.r_init, self.states_init
		out_tensor = self.out_tensor
		trainable_weights = self.trainable_weights
		self.q_init, self.r_init, self.states_init = None, None, None
		self.out_tensor = None
		self.trainable_weights = []
		return q_init, r_init, states_init, out_tensor, trainable_weights

		def set_tensors(self, tensor):
		self.q_init, self.r_init, self.states_init, self.out_tensor, self.trainable_weights = tensor


		class IterRefLSTMEmbedding(Layer):
		"""Embeds its inputs using an LSTM layer."""

		@@ -1631,7 +1648,7 @@ class IterRefLSTMEmbedding(Layer):

		self.trainable_weights = []

		def create_tensor(self, in_layers=None, set_tensors=True):
		def create_tensor(self, in_layers=None, **kwargs):
		"""Execute this layer on input tensors.

		Parameters
		@@ -1687,6 +1704,18 @@ class IterRefLSTMEmbedding(Layer):

		return [x + p, xp + q]

		def none_tensors(self):
		q_init, r_init, states_init = self.q_init, self.r_init, self.states_init
		out_tensor = self.out_tensor
		trainable_weights = self.trainable_weights
		self.q_init, self.r_init, self.states_init = None, None, None
		self.out_tensor = None
		self.trainable_weights = []
		return q_init, r_init, states_init, out_tensor, trainable_weights

		def set_tensors(self, tensor):
		self.q_init, self.r_init, self.states_init, self.out_tensor, self.trainable_weights = tensor


		class BatchNorm(Layer):

deepchem/models/tensorgraph/tests/test_layers_pickle.py

+19 −1

Original line number	Diff line number	Diff line
		@@ -5,7 +5,7 @@ from deepchem.models.tensorgraph.layers import Feature, Conv1D, Dense, Flatten,
		CombineMeanStd, Repeat, GRU, L2Loss, Concat, SoftMax, Constant, Variable, Add, Multiply, InteratomicL2Distances, \
		SoftMaxCrossEntropy, ReduceMean, ToFloat, ReduceSquareDifference, Conv2D, MaxPool, ReduceSum, GraphConv, GraphPool, \
		GraphGather, BatchNorm, WeightedError, \
		LSTMStep
		LSTMStep, AttnLSTMEmbedding
		from deepchem.models.tensorgraph.graph_layers import Combine_AP, Separate_AP, \
		WeaveLayer, WeaveGather, DTNNEmbedding, DTNNGather, DTNNStep, \
		DTNNExtract, DAGLayer, DAGGather, MessagePassing, SetGather
		@@ -442,6 +442,24 @@ def test_MP_pickle():
		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_LSTMStep_pickle():
		"""Tests that LSTMStep can be pickled."""
		n_feat = 10

Admin message