else: os.environ["CUDA_VISIBLE_DEVICES"] = '' config=tf.ConfigProto() config.gpu_options.allow_growth=True sess=tf.Session(config=config) input_i = tf.placeholder(tf.float32, shape=[None,None,None,3]) input_target = tf.placeholder(tf.float32, shape=[None,None,None,3]) lossDict = {} objDict={} with tf.variable_scope(tf.get_variable_scope()): frame_input = input_i[:, :, :, 0:3] frame_processed = input_target[:, :, :, 0:3] with tf.variable_scope('individual'): if not with_IRT: frame_out = net.VCN(frame_input, reuse=False) lossDict["total"] = Lp_loss(frame_out, frame_processed) else: net_pred=net.VCN(frame_input, output_channel=6, reuse=False) frame_out = net_pred[...,0:3] frame_out_minor = net_pred[...,3:6] diff_map = tf.reduce_max(tf.abs(frame_out - frame_processed) / (frame_input + 1e-1), axis=3, keep_dims=True) diff_map1 = tf.reduce_max(tf.abs(frame_out_minor - frame_processed) / (frame_input + 1e-1), axis=3, keep_dims=True) confidence_map = tf.tile(tf.cast(tf.less(diff_map, diff_map1), tf.float32), [1,1,1,3]) lossDict["total"] = tf.reduce_mean(tf.abs(frame_out - frame_processed)) + tf.reduce_mean(tf.abs(frame_out_minor - frame_processed)) lossDict["confidence_map"] = tf.reduce_mean(tf.abs(frame_out - frame_processed)*confidence_map) + tf.reduce_mean( tf.abs(frame_out_minor - frame_processed) * (-confidence_map + 1.0)) opt=tf.train.AdamOptimizer(learning_rate=0.0001).minimize(lossDict["total"], var_list=[var for var in tf.trainable_variables()]) opt_IRT=tf.train.AdamOptimizer(learning_rate=0.0001).minimize(lossDict["confidence_map"], var_list=[var for var in tf.trainable_variables()]) print([var for var in tf.trainable_variables() if var.name.startswith('VCRN')])
def get_model(input): if ARGS.network == "unet": return net.VCN(input, reuse=False)