variable filter sizes enables in UNet (fe97174a) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/models/unet.py

+31 −22

Original line number	Diff line number	Diff line
		@@ -22,28 +22,33 @@ class UNet(TensorGraph):
		number of rows of the image.
		img_cols: int
		number of columns of the image
		filters: list of input
		List of 5 integers to be given which will be the size of the filters in
		each conv layer.
		"""

		def __init__(self,
		img_rows=512,
		img_cols=512,
		filters=[64, 128, 256, 512, 1024],
		model=dc.models.TensorGraph(),
		**kwargs):
		super(UNet, self).__init__(use_queue=False, **kwargs)
		self.img_cols = img_cols
		self.img_rows = img_rows
		self.filters = filters
		self.model = dc.models.TensorGraph()

		input = Feature(shape=(None, self.img_rows, self.img_cols))

		conv1 = Conv2D(
		num_outputs=64,
		num_outputs=self.filters[0],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[input])
		conv1 = Conv2D(
		num_outputs=64,
		num_outputs=self.filters[0],
		kernel_size=3,
		activation='relu',
		padding='same',
		@@ -51,13 +56,13 @@ class UNet(TensorGraph):
		pool1 = MaxPool2D(ksize=2, in_layers=[conv1])

		conv2 = Conv2D(
		num_outputs=128,
		num_outputs=self.filters[1],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[pool1])
		conv2 = Conv2D(
		num_outputs=128,
		num_outputs=self.filters[1],
		kernel_size=3,
		activation='relu',
		padding='same',
		@@ -65,13 +70,13 @@ class UNet(TensorGraph):
		pool2 = MaxPool2D(ksize=2, in_layers=[conv2])

		conv3 = Conv2D(
		num_outputs=256,
		num_outputs=self.filters[2],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[pool2])
		conv3 = Conv2D(
		num_outputs=256,
		num_outputs=self.filters[2],
		kernel_size=3,
		activation='relu',
		padding='same',
		@@ -79,13 +84,13 @@ class UNet(TensorGraph):
		pool3 = MaxPool2D(ksize=2, in_layers=[conv3])

		conv4 = Conv2D(
		num_outputs=512,
		num_outputs=self.filters[3],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[pool3])
		conv4 = Conv2D(
		num_outputs=512,
		num_outputs=self.filters[3],
		kernel_size=3,
		activation='relu',
		padding='same',
		@@ -93,73 +98,77 @@ class UNet(TensorGraph):
		pool4 = MaxPool2D(ksize=2, in_layers=[conv4])

		conv5 = Conv2D(
		num_outputs=1024,
		num_outputs=self.filters[4],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[pool4])
		conv5 = Conv2D(
		num_outputs=1024,
		num_outputs=self.filters[4],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[conv5])

		up6 = Conv2DTranspose(num_outputs=512, kernel_size=2, in_layers=[conv5])
		up6 = Conv2DTranspose(
		num_outputs=self.filters[3], kernel_size=2, in_layers=[conv5])
		concat6 = Concat(in_layers=[conv4, up6], axis=1)
		conv6 = Conv2D(
		num_outputs=512,
		num_outputs=self.filters[3],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[concat6])
		conv6 = Conv2D(
		num_outputs=512,
		num_outputs=self.filters[3],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[conv6])

		up7 = Conv2DTranspose(num_outputs=256, kernel_size=2, in_layers=[conv6])
		up7 = Conv2DTranspose(
		num_outputs=self.filters[2], kernel_size=2, in_layers=[conv6])
		concat7 = Concat(in_layers=[conv3, up7], axis=1)
		conv7 = Conv2D(
		num_outputs=256,
		num_outputs=self.filters[2],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[concat7])
		conv7 = Conv2D(
		num_outputs=256,
		num_outputs=self.filters[2],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[conv7])

		up8 = Conv2DTranspose(num_outputs=128, kernel_size=2, in_layers=[conv7])
		up8 = Conv2DTranspose(
		num_outputs=self.filters[1], kernel_size=2, in_layers=[conv7])
		concat8 = Concat(in_layers=[conv2, up8], axis=1)
		conv8 = Conv2D(
		num_outputs=128,
		num_outputs=self.filters[1],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[concat8])
		conv8 = Conv2D(
		num_outputs=128,
		num_outputs=self.filters[1],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[conv8])

		up9 = Conv2DTranspose(num_outputs=64, kernel_size=2, in_layers=[conv8])
		up9 = Conv2DTranspose(
		num_outputs=self.filters[0], kernel_size=2, in_layers=[conv8])
		concat9 = Concat(in_layers=[conv1, up9], axis=1)
		conv9 = Conv2D(
		num_outputs=64,
		num_outputs=self.filters[0],
		kernel_size=3,
		activation='relu',
		padding='same',
		in_layers=[concat9])
		conv9 = Conv2D(
		num_outputs=64,
		num_outputs=self.filters[0],
		kernel_size=3,
		activation='relu',
		padding='same',

Admin message