yapf (ed720894) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/tests/test_layers_pickle.py

+242 −162

Original line number	Diff line number	Diff line
		@@ -3,7 +3,8 @@ import tensorflow as tf
		from deepchem.models import TensorGraph
		from deepchem.models.tensorgraph.layers import Feature, Conv1D, Dense, Flatten, Reshape, Squeeze, Transpose, \
		CombineMeanStd, Repeat, GRU, L2Loss, Concat, SoftMax, Constant, Variable, Add, Multiply, InteratomicL2Distances, \
		SoftMaxCrossEntropy, ReduceMean, ToFloat, ReduceSquareDifference, Conv2D, MaxPool, ReduceSum
		SoftMaxCrossEntropy, ReduceMean, ToFloat, ReduceSquareDifference, Conv2D, MaxPool, ReduceSum, GraphConv, GraphPool, \
		GraphGather, BatchNorm, WeightedError


		def test_Conv1D_pickle():
		@@ -159,7 +160,11 @@ def testInteratomicL2Distances():
		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])
		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()
		@@ -225,6 +230,7 @@ def test_Conv2D_pickle():
		tg.build()
		tg.save()


		def test_MaxPool_pickle():
		tg = TensorGraph()
		feature = Feature(shape=(tg.batch_size, 10, 10, 10))
		@@ -233,3 +239,77 @@ def test_MaxPool_pickle():
		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 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_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_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()

Admin message