A new deep distortion convolutional neural network for semantic

import logging import os import pickle import time import itertools import tensorflow as tf from tensorflow.contrib.losses.python.metric_learning import triplet_semihard_loss import numpy as np from .util import AttrDict from .util import tf_config from . import util from . import spherical from . import tfnp_compatibility as tfnp from . import params from .layers import * # !!! from . import datasets # !!! logger = logging.getLogger('logger') def dup(x): """ Return two references for input; useful when creating NNs and storing references to layers """ return [x, x] def init_block(args, dset=None): net = {} if dset is None: net['input'], curr = dup(tf.placeholder(args.dtype, shape=(None, *get_indim(args)[1:]))) net['label'] = tf.placeholder('int64', shape=[None]) else: # dset is tuple (iterator, init_ops); iterator returns input and label net['input'], net['label'] = dset[0].get_next() curr = net['input'] net['training'] = tf.placeholder('bool', shape=(), name='is_training') return net, curr def init_sphcnn(args): method = args.transform_method real = args.real_inputs if method == 'naive': fun = lambda *args, **kwargs: spherical.sph_harm_all(*args, **kwargs, real=real) with tf.name_scope('harmonics_or_legendre'): res = args.input_res harmonics = [fun(res // (2**i), as_tfvar=True) for i in range(sum(args.pool_layers) + 1)] return harmonics def two_branch(args, convfun=sphconv, **kwargs): """ Model, that splits input in two branches, and concatenate intermediate feature maps. """ method = args.transform_method l_or_h = init_sphcnn(args) net, curr = init_block(args, **kwargs) assert tfnp.shape(curr)[-1] == 2 curr = [curr[..., 0][..., np.newaxis], curr[..., 1][..., np.newaxis]] # indices for legendre or harmonics high = 0 low = 1 for i, (nf, pool, concat) in enumerate(zip(args.nfilters, args.pool_layers, args.concat_branches)): for b in [0, 1]: name = 'conv{}_b{}'.format(i, b) if concat and b == 0: # top branch also receives features from bottom branch curr[b] = tf.concat(curr, axis=-1) if not pool: with tf.variable_scope(name): net[name], curr[b] = dup(block(args, convfun, net['training'], curr[b], nf, n_filter_params=args.n_filter_params, harmonics_or_legendre=l_or_h[high], method=method)) else: with tf.variable_scope(name): # force spectral pool in first layer if spectral input spectral_pool = True if (args.spectral_input and i == 0) else args.spectral_pool net[name], curr[b] = dup(block(args, convfun, net['training'], curr[b], nf, n_filter_params=args.n_filter_params, harmonics_or_legendre=l_or_h[high], method=method, spectral_pool=pool if spectral_pool else 0, harmonics_or_legendre_low=l_or_h[low])) if not spectral_pool: # weighted avg pooling if args.pool == 'wap': curr[b] = area_weights(tf.layers.average_pooling2d(area_weights(curr[b]), 2*pool, 2*pool, 'same'), invert=True) elif args.pool == 'avg': curr[b] = tf.layers.average_pooling2d(curr[b], 2*pool, 2*pool, 'same') elif args.pool == 'max': curr[b] = tf.layers.max_pooling2d(curr[b], 2*pool, 2*pool, 'same') else: raise ValueError('args.pool') if pool: high += 1 low += 1 # combine for final layer curr = tf.concat(curr, axis=-1) return sphcnn_afterconv(curr, net, args, l_or_h[high]) def sphcnn_afterconv(curr, net, args, l_or_h): """ Part of model after convolutional layers; should be common for different architectures. """ # normalize by area before computing the mean with tf.name_scope('wsa'): if args.weighted_sph_avg: n = tfnp.shape(curr)[1] phi, theta = util.sph_sample(n) phi += np.diff(phi)[0]/2 curr *= np.sin(phi)[np.newaxis, np.newaxis, :, np.newaxis] net['final_conv'] = curr if 'complex' in args.model: curr = tf.abs(curr) nlin = 'relu' else: nlin = args.nonlin # curr is last conv layer with tf.name_scope('final_pool'): net['gap'] = tf.reduce_mean(curr, axis=(1, 2)) if args.final_pool in ['max', 'all']: net['max'] = tf.reduce_max(curr, axis=(1, 2)) if args.final_pool in ['magnitudes', 'all']: net['final_coeffs'] = spherical.sph_harm_transform_batch(curr, method=args.transform_method, harmonics=l_or_h, m0_only=False) # use per frequency magnitudes net['magnitudes'] = tf.contrib.layers.flatten(tf.reduce_sum(tf.square(net['final_coeffs']), axis=(1, 3))) net['magnitudes'] = tf.real(net['magnitudes']) if args.final_pool != 'all': curr = net[args.final_pool] else: curr = tf.concat([net['gap'], net['max'], net['magnitudes']], axis=-1) if args.dropout: curr = tf.nn.dropout(curr, keep_prob=tf.cond(net['training'], lambda: 0.5, lambda: 1.0)) if not args.no_final_fc: with tf.variable_scope('fc1') as scope: net['fc1'], curr = dup(block(AttrDict({**args.__dict__, 'batch_norm': False, 'nonlin': nlin}), tf.layers.dense, net['training'], curr, 64)) if args.dropout: curr = tf.nn.dropout(curr, keep_prob=tf.cond(net['training'], lambda: 0.5, lambda: 1.0)) for v in scope.trainable_variables(): tf.summary.histogram(v.name, v) net['descriptor'] = curr if args.triplet_loss: norm_desc = tf.nn.l2_normalize(curr, dim=-1) # this only works w/ fixed batch size triplet_loss = triplet_semihard_loss(tf.cast(tf.reshape(net['label'], (args.train_bsize,)), 'int32'), norm_desc) # NaNs may appear if bsize is small: triplet_