Changing network layers to handle multiple inputs at once

Aggrathon · Aggrathon · commit 927f81075887 · 2017-04-29T01:34:44.000+03:00
diff --git a/image.py b/image.py
@@ -58,7 +58,7 @@ def start_threads(self):
             t.start()
         if(self.pools > 1):
             print('Processing input images')
-        self.pool = [[self.queue.get() for _ in range(self.batch_size)] for _ in range(self.pools)]
+        self.pool = [[] for _ in range(self.pools)]
 
     def stop_threads(self):
         """Stop the threads that are generating image variations (freeing memory)"""
@@ -67,14 +67,16 @@ def stop_threads(self):
 
     def get_batch(self):
         """Get a batch of images as arrays"""
-        if self.closing:
+        if self.closing:    #Start threads
             self.start_threads()
         self.event.set()
+        if len(self.pool[self.pool_index]) == 0:    #Check and fill image pool
+            self.pool[self.pool_index] = [self.queue.get() for _ in range(self.batch_size)]
         images = self.pool[self.pool_index]
-        for i in range(self.pool_renew):
+        for i in range(self.pool_renew):    #Replace old images
             self.pool[self.pool_index][(self.pool_iteration+i)%self.batch_size] = self.queue.get()
         self.pool_index += 1
-        if self.pool_index == self.pools:
+        if self.pool_index == self.pools:   #Cycle indexes
             self.pool_index = 0
             self.pool_iteration = (self.pool_iteration+self.pool_renew)%self.batch_size
         self.event.clear()
diff --git a/network.py b/network.py
@@ -56,7 +56,7 @@ def __init__(self, name, setup=True, image_size=64, colors=3, batch_size=64, dir
         self._dis_conv = discriminator_convolutions
         self._dis_width = discriminator_base_width
         self._class_depth = classification_depth
-        self.dropout = dropout
+        self._dropout = dropout
         #Training variables
         self.learning_rate = (learning_rate, learning_momentum, learning_momentum2)
         self._y_offset = y_offset
@@ -77,22 +77,21 @@ def __init__(self, name, setup=True, image_size=64, colors=3, batch_size=64, dir
             self.image_input_scaled = tf.subtract(tf.to_float(self.image_input)/127.5, 1, name='image_scaling')
         self.generator_output = None
         self.image_output = self.image_grid_output = None
-        self.image_logit = self.generated_logit = None
         self.variation_updater = self.image_variation = None
         self.generator_solver = self.discriminator_solver = None
         if setup:
             self.setup_network()
 
     def setup_network(self):
         """Initialize the network if it is not done in the constructor"""
-        self.__generator__()
-        self.setup_output()
-        self.discriminator(self.generator_output, self._dis_conv, self._dis_width, self._class_depth)
-        g_loss, d_loss, d_loss_real, d_loss_fake = self.loss_functions(self.image_logit, self.generated_logit, self._y_offset)
+        self.generator_output = self.__generator__([self.generator_input])[0]
+        self.__output__()
+        gen_logit, image_logit = self.__discriminator__([self.generator_output, self.image_input_scaled])
+        g_loss, d_loss, d_loss_real, d_loss_fake = self.loss_functions(image_logit, gen_logit, self._y_offset)
         self.generator_solver, self.discriminator_solver = self.solver_functions(g_loss, d_loss, *self.learning_rate)
 
 
-    def __generator__(self):
+    def __generator__(self, input_tensors):
         """Create a Generator Network"""
         conv_layers = self._gen_conv
         conv_size = self._gen_width
@@ -101,13 +100,15 @@ def __generator__(self):
             conv_image_size = self.image_size // (2**conv_layers)
             assert conv_image_size*(2**conv_layers) == self.image_size, "Images must be a multiple of two (or at least divisible by 2**num_of_conv_layers_plus_one)"
             #Input Layer
-            prev_layer = expand_relu(self.generator_input, [-1, conv_image_size, conv_image_size, conv_size*2**(conv_layers-1)], 'expand')
+            prev_layer = expand_relu(input_tensors, [-1, conv_image_size, conv_image_size, conv_size*2**(conv_layers-1)], 'expand')
             #Conv layers
             for i in range(conv_layers-1):
                 prev_layer = conv2d_transpose(prev_layer, self.batch_size, 2**(conv_layers-i-2)*conv_size, 'convolution_%d'%i)
-            self.generator_output = conv2d_transpose_tanh(prev_layer, self.batch_size, self.colors, 'output')
+            return conv2d_transpose_tanh(prev_layer, self.batch_size, self.colors, 'output')
 
-    def setup_output(self):
+    def __output__(self):
+        if self.log:
+            self.__variation_summary__()
         with tf.name_scope('output'):
             with tf.name_scope("image_list") as scope:
                 self.image_output = tf.cast((self.generator_output + 1) * 127.5, tf.uint8, name=scope)
@@ -129,41 +130,40 @@ def setup_output(self):
                 self.image_grid_output = tf.cast(grid, tf.uint8, name=scope)
 
 
-    def discriminator(self, generator_output, conv_layers, conv_size, class_layers):
+    def __discriminator__(self, input_tensors):
         """Create a Discriminator Network"""
+        conv_layers, conv_size, class_layers = self._dis_conv, self._dis_width, self._class_depth
         image_size = self.image_size
         with tf.variable_scope('discriminator') as scope:
             conv_output_size = ((image_size//(2**conv_layers))**2) * conv_size * conv_layers
             class_output_size = 2**int(math.log(conv_output_size//2, 2))
             #Create Layers
-            def create_network(layer, summary=True):
-                #Convolutional layers
-                for i in range(conv_layers):
-                    layer = conv2d(layer, conv_size*(i+1), name='convolution_%d'%i, norm=(i != 0), summary=summary)
-                layer = tf.reshape(layer, [-1, conv_output_size])
-                #Classification layers
-                for i in range(class_layers):
-                    layer = relu_dropout(layer, class_output_size, self.dropout, 'classification_%d'%i, summary=summary)
-                return linear(layer, 1, 'output', summary=summary)
-            self.generated_logit = create_network(generator_output)
-            scope.reuse_variables()
-            self.image_logit = create_network(self.image_input_scaled, False)
-        if self.log:
-            with tf.variable_scope('pixel_variation'):
-                #Pixel Variations
-                img_tot_var = tf.image.total_variation(self.image_input_scaled)
-                gen_tot_var = tf.image.total_variation(generator_output)
-                image_variation = tf.reduce_sum(img_tot_var)
-                gener_variation = tf.reduce_sum(gen_tot_var)
-                tf.summary.histogram('images', img_tot_var)
-                tf.summary.histogram('generated', gen_tot_var)
-                tf.summary.scalar('images', image_variation)
-                tf.summary.scalar('generated', gener_variation)
-                ema = tf.train.ExponentialMovingAverage(decay=0.95, num_updates=self.iterations)
-                ema_apply = ema.apply([image_variation])
-                self.variation_updater = tf.group(ema_apply)
-                self.image_variation = ema.average(image_variation)
-                tf.summary.scalar('images_averaged', self.image_variation)
+            prev_layer = input_tensors
+            for i in range(conv_layers): #Convolutional layers
+                prev_layer = conv2d(prev_layer, conv_size*(i+1), name='convolution_%d'%i, norm=(i != 0))
+            prev_layer = [tf.reshape(layer, [-1, conv_output_size]) for layer in prev_layer]
+            for i in range(class_layers): #Classification layers
+                prev_layer = relu_dropout(prev_layer, class_output_size, self._dropout, 'classification_%d'%i)
+            prev_layer = linear(prev_layer, 1, 'output')
+        return prev_layer
+
+    def __variation_summary__(self):
+        """Create summaries for pixel variation"""
+        with tf.variable_scope('pixel_variation'):
+            #Pixel Variations
+            img_tot_var = tf.image.total_variation(self.image_input_scaled)
+            gen_tot_var = tf.image.total_variation(self.generator_output)
+            image_variation = tf.reduce_sum(img_tot_var)
+            gener_variation = tf.reduce_sum(gen_tot_var)
+            tf.summary.histogram('images', img_tot_var)
+            tf.summary.histogram('generated', gen_tot_var)
+            tf.summary.scalar('images', image_variation)
+            tf.summary.scalar('generated', gener_variation)
+            ema = tf.train.ExponentialMovingAverage(decay=0.95, num_updates=self.iterations)
+            ema_apply = ema.apply([image_variation])
+            self.variation_updater = tf.group(ema_apply)
+            self.image_variation = ema.average(image_variation)
+            tf.summary.scalar('images_averaged', self.image_variation)
 
 
     def loss_functions(self, real_logit, fake_logit, y_offset=0):
diff --git a/operators.py b/operators.py
@@ -43,65 +43,80 @@ def lrelu(tensor, leak: float=0.2):
 
 # Network Layers
 
-def conv2d(tensor, output_size: int, name: str='conv2d', norm: bool=True, stddev: float=0.02, term: float=0.01, summary: bool=True):
+def conv2d(tensors, output_size: int, name: str='conv2d', norm: bool=True, stddev: float=0.02, term: float=0.01, summary: bool=True):
 	"""Create a convolutional layer"""
 	with tf.variable_scope(name):
-		weight, bias = weight_bias([5, 5, int(tensor.get_shape()[-1]), output_size], stddev, term, summary)
-		conv = tf.nn.conv2d(tensor, weight, [1, 2, 2, 1], "SAME")
-		if norm:
-			conv = tf.contrib.layers.batch_norm(conv, decay=0.9, updates_collections=None, scale=False,
-				trainable=True, reuse=True, scope="normalization", is_training=True, epsilon=0.00001)
-		return lrelu(tf.nn.bias_add(conv, bias))
+		weight, bias = weight_bias([5, 5, int(tensors[0].get_shape()[-1]), output_size], stddev, term, summary)
+		output = []
+		for tensor in tensors:
+			conv = tf.nn.conv2d(tensor, weight, [1, 2, 2, 1], "SAME")
+			if norm:
+				conv = tf.contrib.layers.batch_norm(conv, decay=0.9, updates_collections=None, scale=False,
+					trainable=True, reuse=True, scope="normalization", is_training=True, epsilon=0.00001)
+			output.append(lrelu(tf.nn.bias_add(conv, bias)))
+		return output
 
 def relu(tensor, output_size: int, name: str='relu', stddev: float=0.02, term: float=0.01, summary: bool=True):
 	"""Create a relu layer"""
 	with tf.variable_scope(name):
 		weight, bias = weight_bias([int(tensor.get_shape()[-1]), output_size], stddev, term, summary)
 		return tf.nn.relu(tf.matmul(tensor, weight) + bias)
 
-def relu_dropout(tensor, output_size: int, dropout: float=0.4, name: str='relu_dropout', stddev: float=0.02, term: float=0.01, summary: bool=True):
+def relu_dropout(tensors, output_size: int, dropout: float=0.4, name: str='relu_dropout', stddev: float=0.02, term: float=0.01, summary: bool=True):
 	"""Create a relu layer with dropout"""
 	with tf.variable_scope(name):
-		weight, bias = weight_bias([int(tensor.get_shape()[-1]), output_size], stddev, term, summary)
-		relu_layer = tf.nn.relu(tf.matmul(tensor, weight) + bias)
-		return tf.nn.dropout(relu_layer, dropout)
-
-def linear(tensor, output_size: int, name: str='linear', stddev: float=0.02, term: float=0.01, summary: bool=True):
+		weight, bias = weight_bias([int(tensors[0].get_shape()[-1]), output_size], stddev, term, summary)
+		output = []
+		for tensor in tensors:
+			relu_layer = tf.nn.relu(tf.matmul(tensor, weight) + bias)
+			output.append(tf.nn.dropout(relu_layer, dropout))
+		return output
+
+def linear(tensors, output_size: int, name: str='linear', stddev: float=0.02, term: float=0.01, summary: bool=True):
 	'''Create a fully connected layer'''
 	with tf.variable_scope(name):
-		weight, bias = weight_bias([tensor.get_shape()[-1], output_size], stddev, term, summary)
-		return tf.matmul(tensor, weight) + bias
+		weight, bias = weight_bias([tensors[0].get_shape()[-1], output_size], stddev, term, summary)
+		return [tf.matmul(tensor, weight) + bias for tensor in tensors]
 
-def conv2d_transpose(tensor, batch_size=1, conv_size=32, name: str='conv2d_transpose', norm: bool=True, stddev: float=0.02, term: float=0.01, summary: bool=True):
+def conv2d_transpose(tensors, batch_size=1, conv_size=32, name: str='conv2d_transpose', norm: bool=True, stddev: float=0.02, term: float=0.01, summary: bool=True):
 	"""Create a transpose convolutional layer"""
 	with tf.variable_scope(name):
-		tensor_shape = tensor.get_shape()
+		tensor_shape = tensors[0].get_shape()
 		filt, bias = filter_bias([5, 5, conv_size, tensor_shape[-1]], stddev, term, summary)
 		conv_shape = [batch_size, int(tensor_shape[1]*2), int(tensor_shape[2]*2), conv_size]
-		deconv = tf.nn.conv2d_transpose(tensor, filt, conv_shape, [1, 2, 2, 1])
-		if norm:
-			deconv = tf.contrib.layers.batch_norm(deconv, decay=0.9, updates_collections=None, scale=False,
-				trainable=True, reuse=True, scope="normalization", is_training=True, epsilon=0.00001)
-		return tf.nn.relu(tf.nn.bias_add(deconv, bias))
-
-def conv2d_transpose_tanh(tensor, batch_size=1, conv_size=32, name: str='conv2d_transpose_tanh', stddev: float=0.02, summary: bool=True):
+		output = []
+		for tensor in tensors:
+			deconv = tf.nn.conv2d_transpose(tensor, filt, conv_shape, [1, 2, 2, 1])
+			if norm:
+				deconv = tf.contrib.layers.batch_norm(deconv, decay=0.9, updates_collections=None, scale=False,
+					trainable=True, reuse=True, scope="normalization", is_training=True, epsilon=0.00001)
+			output.append(tf.nn.relu(tf.nn.bias_add(deconv, bias)))
+		return output
+
+def conv2d_transpose_tanh(tensors, batch_size=1, conv_size=32, name: str='conv2d_transpose_tanh', stddev: float=0.02, summary: bool=True):
 	"""Create a transpose convolutional layer"""
 	with tf.variable_scope(name):
-		tensor_shape = tensor.get_shape()
+		tensor_shape = tensors[0].get_shape()
 		filt = tf.get_variable('filter', [5, 5, conv_size, tensor_shape[-1]], tf.float32, tf.random_normal_initializer(0, stddev), trainable=True)
-		conv_shape = [batch_size, int(tensor_shape[1]*2), int(tensor_shape[2]*2), conv_size]
-		deconv = tf.nn.conv2d_transpose(tensor, filt, conv_shape, [1, 2, 2, 1])
-		return tf.nn.tanh(deconv)
-
-def expand_relu(tensor, out_shape, name: str='expand_relu', norm: bool=True, stddev: float=0.2, term: float=0.01, summary: bool=True):
+		output = []
+		for tensor in tensors:
+			conv_shape = [batch_size, int(tensor_shape[1]*2), int(tensor_shape[2]*2), conv_size]
+			deconv = tf.nn.conv2d_transpose(tensor, filt, conv_shape, [1, 2, 2, 1])
+			output.append(tf.nn.tanh(deconv))
+		return output
+
+def expand_relu(tensors, out_shape, name: str='expand_relu', norm: bool=True, stddev: float=0.2, term: float=0.01, summary: bool=True):
 	"""Create a layer that expands an input to a shape"""
 	with tf.variable_scope(name) as scope:
-		weight, bias = weight_bias([tensor.get_shape()[-1], np.prod(out_shape[1:])], stddev, term, summary)
-		lin = tf.matmul(tensor, weight) + bias
-		reshape = tf.reshape(lin, out_shape)
-		if norm:
-			reshape = tf.contrib.layers.batch_norm(reshape, decay=0.9, updates_collections=None, scale=False,
-				trainable=True, reuse=True, scope=scope, is_training=True, epsilon=0.00001)
-		return tf.nn.relu(reshape)
+		weight, bias = weight_bias([tensors[0].get_shape()[-1], np.prod(out_shape[1:])], stddev, term, summary)
+		output = []
+		for tensor in tensors:
+			lin = tf.matmul(tensor, weight) + bias
+			reshape = tf.reshape(lin, out_shape)
+			if norm:
+				reshape = tf.contrib.layers.batch_norm(reshape, decay=0.9, updates_collections=None, scale=False,
+					trainable=True, reuse=True, scope=scope, is_training=True, epsilon=0.00001)
+			output.append(tf.nn.relu(reshape))
+		return output
 
 
