Upload net.py

net.py

@@ -31,22 +31,48 @@ def data_loader_bd_rm_from_tfrecord(batch_size=1):
 	return loader_dict, num_batch
 
 class Network(object):
-    def __init__(self, dtype=tf.float32):
-        print('Initial nn network object...')
-        self.dtype = dtype
-        # ... existing code ...
-
-    def convert_one_hot_to_image(self, one_hot, dtype='float', act=None):
-        import tensorflow.compat.v1 as tf
-        if act == 'softmax':
-            one_hot = tf.nn.softmax(one_hot, axis=-1)
-        [n, h, w, c] = one_hot.shape.as_list()
-        im = tf.reshape(tf.argmax(one_hot, axis=-1), [n, h, w, 1])
-        if dtype == 'int':
-            im = tf.cast(im, dtype=tf.uint8)
-        else:
-            im = tf.cast(im, dtype=tf.float32)
-        return im
+	"""docstring for Network"""
+	def __init__(self, dtype=tf.float32):
+		print('Initial nn network object...')
+		self.dtype = dtype
+		self.pre_train_restore_map = {'vgg_16/conv1/conv1_1/weights':'FNet/conv1_1/W', # {'checkpoint_scope_var_name':'current_scope_var_name'} shape must be the same
+									'vgg_16/conv1/conv1_1/biases':'FNet/conv1_1/b',	
+									'vgg_16/conv1/conv1_2/weights':'FNet/conv1_2/W',
+									'vgg_16/conv1/conv1_2/biases':'FNet/conv1_2/b',	
+									'vgg_16/conv2/conv2_1/weights':'FNet/conv2_1/W',
+									'vgg_16/conv2/conv2_1/biases':'FNet/conv2_1/b',	
+									'vgg_16/conv2/conv2_2/weights':'FNet/conv2_2/W',
+									'vgg_16/conv2/conv2_2/biases':'FNet/conv2_2/b',	
+									'vgg_16/conv3/conv3_1/weights':'FNet/conv3_1/W',
+									'vgg_16/conv3/conv3_1/biases':'FNet/conv3_1/b',	
+									'vgg_16/conv3/conv3_2/weights':'FNet/conv3_2/W',
+									'vgg_16/conv3/conv3_2/biases':'FNet/conv3_2/b',	
+									'vgg_16/conv3/conv3_3/weights':'FNet/conv3_3/W',
+									'vgg_16/conv3/conv3_3/biases':'FNet/conv3_3/b',	
+									'vgg_16/conv4/conv4_1/weights':'FNet/conv4_1/W',
+									'vgg_16/conv4/conv4_1/biases':'FNet/conv4_1/b',	
+									'vgg_16/conv4/conv4_2/weights':'FNet/conv4_2/W',
+									'vgg_16/conv4/conv4_2/biases':'FNet/conv4_2/b',	
+									'vgg_16/conv4/conv4_3/weights':'FNet/conv4_3/W',
+									'vgg_16/conv4/conv4_3/biases':'FNet/conv4_3/b',	
+									'vgg_16/conv5/conv5_1/weights':'FNet/conv5_1/W',
+									'vgg_16/conv5/conv5_1/biases':'FNet/conv5_1/b',	
+									'vgg_16/conv5/conv5_2/weights':'FNet/conv5_2/W',
+									'vgg_16/conv5/conv5_2/biases':'FNet/conv5_2/b',	
+									'vgg_16/conv5/conv5_3/weights':'FNet/conv5_3/W',
+									'vgg_16/conv5/conv5_3/biases':'FNet/conv5_3/b'} 
+
+	def convert_one_hot_to_image(self, one_hot, dtype='float', act=None):
+		# This method was moved from MODEL in main.py for inference compatibility
+		if act == 'softmax':
+			one_hot = tf.nn.softmax(one_hot, axis=-1)
+		[n, h, w, c] = one_hot.shape.as_list()
+		im = tf.reshape(tf.argmax(one_hot, axis=-1), [n, h, w, 1])
+		if dtype == 'int':
+			im = tf.cast(im, dtype=tf.uint8)
+		else:
+			im = tf.cast(im, dtype=tf.float32)
+		return im
 
 	# basic layer 
 	def _he_uniform(self, shape, regularizer=None, trainable=None, name=None):