change input to RGB and update ksize in MaxPool2D (1a200012) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/models/unet.py

+5 −5

Original line number	Diff line number	Diff line
		@@ -39,7 +39,7 @@ class UNet(TensorGraph):
		self.filters = filters
		self.model = dc.models.TensorGraph()

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

		conv1 = Conv2D(
		num_outputs=self.filters[0],
		@@ -53,7 +53,7 @@ class UNet(TensorGraph):
		activation='relu',
		padding='same',
		in_layers=[conv1])
		pool1 = MaxPool2D(ksize=2, in_layers=[conv1])
		pool1 = MaxPool2D(ksize=[1,2,2,3], 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=2, in_layers=[conv2])
		pool2 = MaxPool2D(ksize=[1,2,2,3], 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=2, in_layers=[conv3])
		pool3 = MaxPool2D(ksize=[1,2,2,3], 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=2, in_layers=[conv4])
		pool4 = MaxPool2D(ksize=[1,2,2,3], in_layers=[conv4])

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