normalizing ynet model build methods to fcnet (b1bee656) · Commits · disMultiABM / stemdl

stemdl/losses.py

+9 −8

Original line number	Diff line number	Diff line
		@@ -122,6 +122,7 @@ def calc_loss(n_net, scope, hyper_params, params, labels, step=None, images=None
		losses = [inverter_loss , decoder_loss_re, decoder_loss_im]
		# losses, prefac = ynet_adjusted_losses(losses, step)
		# tf.summary.scalar("prefac_inverter", prefac)
		# losses = [inverter_loss]
		regularization = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
		# Calculate the total loss
		total_loss = tf.add_n(losses + regularization, name='total_loss')
		@@ -227,27 +228,27 @@ def ynet_adjusted_losses(losses, global_step):
		'''
		Schedule the different loss components based on global training step
		'''
		threshold = tf.constant(0.25)
		max_prefac = 20
		threshold = tf.constant(0.8)
		max_prefac = 1
		ema = tf.train.ExponentialMovingAverage(0.9)
		loss_averages_op = ema.apply(losses)
		prefac_initial = 2.0
		prefac_initial = 1e-5

		with tf.control_dependencies([loss_averages_op]):
		def ramp():
		prefac = tf.cast(tf.train.exponential_decay(tf.constant(prefac_initial), global_step, 10000,
		tf.constant(0.5, dtype=tf.float32), staircase=False), tf.float32)
		prefac = tf.cast(tf.train.exponential_decay(tf.constant(prefac_initial), global_step, 1000,
		tf.constant(0.1, dtype=tf.float32), staircase=False), tf.float32)
		prefac = tf.constant(prefac_initial) ** 2 * tf.pow(prefac, tf.constant(-1., dtype=tf.float32))
		prefac = tf.minimum(prefac, tf.cast(max_prefac, tf.float32))
		return prefac

		def decay(prefac_current):
		prefac = tf.train.exponential_decay(prefac_current, global_step, 1000, tf.constant(0.5, dtype=tf.float32),
		prefac = tf.train.exponential_decay(prefac_current, global_step, 1000, tf.constant(0.1, dtype=tf.float32),
		staircase=True)
		return prefac
		inv_loss, dec_re_loss, dec_im_loss = losses

		prefac = tf.cond(inv_loss > threshold, true_fn=ramp, false_fn=lambda: decay(ramp()))
		prefac = tf.cond(inv_loss > threshold, false_fn=ramp, true_fn=lambda: decay(prefac_initial))
		tf.summary.scalar("prefac_inverter", prefac)
		losses = [prefac * (inv_loss - threshold), dec_re_loss, dec_im_loss]
		losses = [inv_loss , prefac * dec_re_loss, prefac * dec_im_loss]
		return losses, prefac

stemdl/network.py

+76 −105

Original line number	Diff line number	Diff line
		@@ -994,7 +994,6 @@ class ConvNet:
		scaled = scaled * (scale[-1] - scale[0]) + scale[0]
		return scaled


		class ResNet(ConvNet):

		# def __init__(self, args, *kwargs):
		@@ -1104,7 +1103,6 @@ class ResNet(ConvNet):
		else:
		super(ResNet, self)._print_layer_specs(params, scope, input_shape, output_shape)


		class FCDenseNet(ConvNet):
		"""
		Fully Convolutional Dense Neural Network
		@@ -1439,7 +1437,6 @@ class FCDenseNet(ConvNet):
		block_features = tf.concat(block_features, axis=1)
		return input, block_features


		class FCNet(ConvNet):
		"""
		Fully (vanilla) convolutional neural net
		@@ -2236,7 +2233,6 @@ class FCNet(ConvNet):

		return kernel


		class YNet(FCDenseNet, FCNet):
		"""
		An inverter model
		@@ -2797,14 +2793,25 @@ class YNet(FCDenseNet, FCNet):
		# if self.summary:
		# self._activation_summary(out)
		# self._activation_image_summary(out)

		with tf.variable_scope('%s_CONV_FIN' % subnet, reuse=self.reuse) as scope:
		conv_1by1 = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [3, 3], 'padding': 'SAME', 'features': 1})
		self.print_verbose(">>> Adding CONV_FIN layer: ")
		self.print_verbose(' input: %s' %format(out.get_shape().as_list()))
		out, _ = self._conv(input=out, params=conv_1by1)
		self.print_verbose(' output: %s' %format(out.get_shape().as_list()))
		out_shape = out.get_shape().as_list()
		self._print_layer_specs(layer_params, scope, in_shape, out_shape)
		self.scopes.append(scope)
		if self.summary:
		self._activation_summary(out)
		self._activation_image_summary(out)
		self.model_output[subnet] = out
		self.update_all_attrs(subnet=subnet)
		self.print_rank('Total # of blocks: %d, weights: %2.1e, memory: %s MB, ops: %3.2e \n' % (len(network),
		self.num_weights,
		format(self.mem / 1024),
		self.get_ops()))
		# if subnet == 'inverter':
		# out = tf.reduce_mean(out, axis=1, keepdims=True)
		self.model_output[subnet] = out
		self.update_all_attrs(subnet=subnet)

		def _upscale(self, inputs=None, params=None, scale=2):
		conv_params = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [1, 1],
		@@ -2841,20 +2848,14 @@ class YNet(FCDenseNet, FCNet):
		scopes_list = []
		out = self.model_output['encoder']
		out_shape = out.shape.as_list()
		params = self.network['encoder']['freq2space']['cvae_params']
		fully_connected = params['fc_params']
		num_fc = params['n_fc_layers']
		conv_1by1_1 = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [1, 1],
		'features': 1,
		'activation': 'relu',
		'padding': 'VALID',
		'batch_norm': True, 'dropout':0.0})
		conv_1by1_1024 = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [1, 1],
		'features': 1024,
		'activation': 'relu',
		params = self.network['encoder']['freq2space']
		fully_connected = params['cvae_params']['fc_params']
		num_fc = params['cvae_params']['n_fc_layers']
		conv_1by1 = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [1, 1],
		'features': params['init_features'],
		'activation': "relu",
		'padding': 'VALID',
		'batch_norm': True, 'dropout':0.0})
		if True:
		'batch_norm': True, 'dropout':0})
		def fc_map(tens):
		for i in range(num_fc):
		with tf.variable_scope('%s_fc_%d' %(subnet, i), reuse=self.reuse) as scope :
		@@ -2866,47 +2867,33 @@ class YNet(FCDenseNet, FCNet):
		out = tf.transpose(out, perm= [1, 2, 0])
		dim = int(math.sqrt(self.images.shape.as_list()[1]))
		out = tf.reshape(out, [self.params['batch_size'], -1, dim, dim])
		# else:
		# out = tf.reshape(out, [out_shape[0]*out_shape[1], out_shape[2], out_shape[3], out_shape[4]])
		# with tf.variable_scope('%s_conv_1by1_1' % subnet, reuse=self.reuse) as scope:
		# out, _ = self._conv(input=out, params=conv_1by1_1)
		# out = tf.reshape(out, [out_shape[0], out_shape[1], out_shape[3], out_shape[4]])
		# out = tf.transpose(out, perm=[1,0,2,3])

		# # scopes_list.append(scope)
		# with tf.variable_scope('%s_conv_1by1_1024' % subnet, reuse=self.reuse) as scope:
		# out, _ = self._conv(input=out, params=conv_1by1_1024)
		# out = self._activate(input=out, params=conv_1by1_1024)
		# do_bn = conv_1by1_1024.get('batch_norm', False)
		# if do_bn:
		# out = self._batch_norm(input=out)
		# else:
		# out = self._add_bias(input=out, params=conv_1by1_1024)
		# # scopes_list.append(scope)
		with tf.variable_scope('%s_conv_1by1' % subnet, reuse=self.reuse) as scope:
		out, _ = self._conv(input=out, params=conv_1by1)
		out = self._activate(input=out, params=conv_1by1)
		do_bn = conv_1by1.get('batch_norm', False)
		if do_bn:
		out = self._batch_norm(input=out)
		else:
		out = self._add_bias(input=out, params=conv_1by1)
		out = self._activate(input=out, params=conv_1by1)
		scopes_list.append(scope)

		self._build_branch(subnet=subnet, inputs=out)

		# conv_1by1 = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [1, 1], 'features': 1,
		# 'activation': 'relu', 'padding': 'SAME', 'batch_norm': False})
		# with tf.variable_scope('%s_CONV_FIN' % subnet, reuse=self.reuse) as scope:
		# out, _ = self._conv(input=self.model_output[subnet], params=conv_1by1)
		# do_bn = conv_1by1.get('batch_norm', False)
		# if do_bn:
		# out = self._batch_norm(input=out)
		# else:
		# out = self._add_bias(input=out, params=conv_1by1)
		# out = self._activate(input=out, params=conv_1by1)
		# scopes_list.append(scope)
		# self.model_output[subnet] = out
		self.all_scopes[subnet] += scopes_list

		def build_inverter(self):
		out = self.model_output['encoder']
		params = self.network['encoder']['freq2space']['cvae_params']
		fully_connected = params['fc_params']
		num_fc = params['n_fc_layers']
		scopes_list = []
		if True:
		out = self.model_output['encoder']
		out_shape = out.shape.as_list()
		params = self.network['encoder']['freq2space']
		fully_connected = params['cvae_params']['fc_params']
		num_fc = params['cvae_params']['n_fc_layers']
		conv_1by1 = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [1, 1],
		'features': params['init_features'],
		'activation': "relu",
		'padding': 'VALID',
		'batch_norm': True, 'dropout':0})
		def fc_map(tens):
		for i in range(num_fc):
		with tf.variable_scope('Inverter_fc_%d' %i, reuse=self.reuse) as scope :
		@@ -2914,37 +2901,21 @@ class YNet(FCDenseNet, FCNet):
		tens = self._activate(input=tens, params=fully_connected)
		scopes_list.append(scope)
		return tens
		out = tf.map_fn(fc_map, out, back_prop=True, swap_memory=True, parallel_iterations=256)
		out = tf.map_fn(fc_map, out, back_prop=True)
		out = tf.transpose(out, perm= [1, 2, 0])
		dim = int(math.sqrt(self.images.shape.as_list()[1]))
		out = tf.reshape(out, [self.params['batch_size'], -1, dim, dim])
		# else:
		# conv_1by1_1 = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [1, 1],
		# 'features': 1,
		# 'activation': 'relu',
		# 'padding': 'VALID',
		# 'batch_norm': True, 'dropout':0.0})
		# out_shape = out.shape.as_list()
		# out = tf.reshape(out, [out_shape[0]*out_shape[1], out_shape[2], out_shape[3], out_shape[4]])
		# with tf.variable_scope('%s_conv_1by1_1' % 'inverter', reuse=self.reuse) as scope:
		# out, _ = self._conv(input=out, params=conv_1by1_1)
		# out = tf.reshape(out, [out_shape[0], out_shape[1], out_shape[3], out_shape[4]])
		# out = tf.transpose(out, perm=[1,0,2,3])
		# scopes_list.append(scope)
		# conv_1by1_1024 = OrderedDict({'type': 'conv_2D', 'stride': [1, 1], 'kernel': [1, 1],
		# 'features': 1024,
		# 'activation': 'relu',
		# 'padding': 'VALID',
		# 'batch_norm': True, 'dropout':0.0})
		# with tf.variable_scope('inverter_conv_1by1_1024', reuse=self.reuse) as scope:
		# out, _ = self._conv(input=out, params=conv_1by1_1024)
		# do_bn = conv_1by1_1024.get('batch_norm', False)
		# if do_bn:
		# out = self._batch_norm(input=out)
		# else:
		# out = self._add_bias(input=out, params=conv_1by1_1024)
		# out = self._activate(input=out, params=conv_1by1_1024)
		# # scopes_list.append(scope)
		with tf.variable_scope('%s_conv_1by1' % 'inverter', reuse=self.reuse) as scope:
		out, _ = self._conv(input=out, params=conv_1by1)
		out = self._activate(input=out, params=conv_1by1)
		do_bn = conv_1by1.get('batch_norm', False)
		if do_bn:
		out = self._batch_norm(input=out)
		else:
		out = self._add_bias(input=out, params=conv_1by1)
		out = self._activate(input=out, params=conv_1by1)
		scopes_list.append(scope)

		self._build_branch(subnet='inverter', inputs=out)
		self.all_scopes['inverter'] += scopes_list

stemdl/runtime.py

+0 −8

Original line number	Diff line number	Diff line
		@@ -35,7 +35,6 @@ def tensorflow_version_tuple():
		major, minor, patch = v.split('.')
		return (int(major), int(minor), patch)


		def float32_variable_storage_getter(getter, name, shape=None, dtype=None,
		initializer=None, regularizer=None,
		trainable=True,
		@@ -49,7 +48,6 @@ def float32_variable_storage_getter(getter, name, shape=None, dtype=None,
		variable = tf.cast(variable, dtype)
		return variable


		class TrainHelper(object):
		def __init__(self, params, saver, writer, net_ops, last_step=0, log_freq=1):
		self.params = params
		@@ -129,12 +127,10 @@ class TrainHelper(object):
		print_rank('loss is nan... Exiting!')
		sys.exit(0)


		def print_rank(args, *kwargs):
		if hvd.rank() == 0:
		print(args, *kwargs)


		def train(network_config, hyper_params, params):
		"""
		Train the network for a number of steps using horovod and asynchronous I/O staging ops.
		@@ -418,8 +414,6 @@ def train(network_config, hyper_params, params):
		print_rank('Saved Final Checkpoint.')
		return



		def train_inverter(network_config, hyper_params, params):
		"""
		Train the network for a number of steps using horovod and asynchronous I/O staging ops.
		@@ -703,7 +697,6 @@ def train_inverter(network_config, hyper_params, params):
		validateStep += params['validate_step']
		#next_validation_epoch += params['epochs_per_validation']


		def validate(network_config, hyper_params, params, sess, dset, num_batches=10):
		"""
		Runs validation with current weights
		@@ -819,7 +812,6 @@ def validate(network_config, hyper_params, params, sess, dset, num_batches=10):
		print_rank('Validation Reconstruction Error %s: %3.3e' % (loss_label, errors.mean()))
		tf.summary.scalar("Validation_loss_label_%s" % loss_label, tf.constant(errors.mean()))


		def validate_ckpt(network_config, hyper_params, params, num_batches=None,
		last_model= False, sleep=-1):
		"""

Admin message