more tests for UNET (d53a23f5) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/models/unet.py

+4 −4

Original line number	Diff line number	Diff line
		@@ -53,7 +53,7 @@ class UNet(TensorGraph):
		activation='relu',
		padding='same',
		in_layers=[conv1])
		pool1 = MaxPool2D(ksize=[1, 2, 2, 3], in_layers=[conv1])
		pool1 = MaxPool2D(ksize=[1, 2, 2, 1], in_layers=[conv1])

		conv2 = Conv2D(
		num_outputs=self.filters[1],
		@@ -67,7 +67,7 @@ class UNet(TensorGraph):
		activation='relu',
		padding='same',
		in_layers=[conv2])
		pool2 = MaxPool2D(ksize=[1, 2, 2, 3], in_layers=[conv2])
		pool2 = MaxPool2D(ksize=[1, 2, 2, 1], in_layers=[conv2])

		conv3 = Conv2D(
		num_outputs=self.filters[2],
		@@ -81,7 +81,7 @@ class UNet(TensorGraph):
		activation='relu',
		padding='same',
		in_layers=[conv3])
		pool3 = MaxPool2D(ksize=[1, 2, 2, 3], in_layers=[conv3])
		pool3 = MaxPool2D(ksize=[1, 2, 2, 1], in_layers=[conv3])

		conv4 = Conv2D(
		num_outputs=self.filters[3],
		@@ -95,7 +95,7 @@ class UNet(TensorGraph):
		activation='relu',
		padding='same',
		in_layers=[conv4])
		pool4 = MaxPool2D(ksize=[1, 2, 2, 3], in_layers=[conv4])
		pool4 = MaxPool2D(ksize=[1, 2, 2, 1], in_layers=[conv4])

		conv5 = Conv2D(
		num_outputs=self.filters[4],

+22 −3

Original line number	Diff line number	Diff line
		import deepchem as dc
		import deepchem
		import numpy as np
		import tensorflow as tf
		import unittest
		@@ -10,8 +10,27 @@ class TestUNet(unittest.TestCase):

		def test_unet(self):
		unet2D = deepchem.models.tensorgraph.models.unet.UNet(
		learning_rate=0.003, img_rows=32, img_cols=32)
		learning_rate=0.003, img_rows=32, img_cols=32, model_dir='./unet/')
		# Prepare Training Data
		data = np.ones((5, 32, 32, 3))
		labels = np.ones((5, 32 * 32))
		train = deepchem.data.NumpyDataset(data, labels)
		unet2D.fit(train, nb_epochs=0)
		# Train the model
		unet2D.fit(train, nb_epochs=2)
		unet2D.save()
		# Prepare the Testing data
		test_data = np.ones((2, 32, 32, 3))
		test = deepchem.data.NumpyDataset(test_data)
		# predict
		predictions = unet2D.predict(test)
		# check output shape
		assert predictions.shape == (2, 32, 32, 1)

		# new object of UNet to test if loading the model results in same predictions
		unet2D_new = deepchem.models.tensorgraph.models.unet.UNet(
		learning_rate=0.003, img_rows=32, img_cols=32, model_dir='./unet/')
		unet2D_new.load_from_dir('./unet/')
		unet2D_new.restore()
		predictions_new = unet2D_new.predict(test)

		assert np.all(predictions == predictions_new)