def prob_unet2D(z_list, x, zdim_0, n_classes, generation_mode, training, scope_reuse=False, norm=tfnorm.batch_norm, **kwargs): resolution_levels = kwargs.get('resolution_levels', 7) n0 = kwargs.get('n0', 32) num_channels = [n0, 2*n0, 4*n0,6*n0, 6*n0, 6*n0, 6*n0] conv_unit = layers.conv2D with tf.variable_scope('prior') as scope: if scope_reuse: scope.reuse_variables() add_bias = False if norm == tfnorm.batch_norm else True enc = [] for ii in range(resolution_levels): enc.append([]) # In first layer set input to x rather than max pooling if ii == 0: enc[ii].append(x) else: enc[ii].append(layers.averagepool2D(enc[ii-1][-1])) enc[ii].append(conv_unit(enc[ii][-1], 'conv_%d_1' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) enc[ii].append(conv_unit(enc[ii][-1], 'conv_%d_2' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) enc[ii].append(conv_unit(enc[ii][-1], 'conv_%d_3' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) mu_p = conv_unit(enc[-1][-1], 'pre_mu', num_filters=zdim_0, kernel_size=(1, 1), activation=tf.identity) mu = [layers.global_averagepool2D(mu_p)] sigma_p = conv_unit(enc[-1][-1], 'pre_sigma', num_filters=zdim_0, kernel_size=(1, 1), activation=tf.nn.softplus) sigma = [layers.global_averagepool2D(sigma_p)] z = [mu[0] + sigma[0] * tf.random_normal(tf.shape(mu[0]), 0, 1, dtype=tf.float32)] return z, mu, sigma
def hybrid(x, s_oh, zdim_0, training, scope_reuse=False, norm=tfnorm.batch_norm, **kwargs): n0 = kwargs.get('n0', 32) num_channels = [n0, 2 * n0, 4 * n0, 6 * n0, 6 * n0, 6 * n0, 6 * n0] with tf.variable_scope('posterior') as scope: if scope_reuse: scope.reuse_variables() full_cov_list = kwargs.get('full_cov_list', None) n0 = kwargs.get('n0', 32) latent_levels = kwargs.get('latent_levels', 5) resolution_levels = kwargs.get('resolution_levels', 7) spatial_xdim = x.get_shape().as_list()[1:3] full_latent_dependencies = kwargs.get('full_latent_dependencies', False) pre_z = [None] * resolution_levels mu = [None] * latent_levels sigma = [None] * latent_levels z = [None] * latent_levels z_ups_mat = [] for i in range(latent_levels): z_ups_mat.append( [None] * latent_levels) # encoding [original resolution][upsampled to] # Generate pre_z's for i in range(resolution_levels): if i == 0: net = tf.concat([x, s_oh - 0.5], axis=-1) else: net = layers.averagepool2D(pre_z[i - 1]) net = layers.conv2D(net, 'z%d_pre_1' % i, num_filters=num_channels[i], normalisation=norm, training=training) net = layers.conv2D(net, 'z%d_pre_2' % i, num_filters=num_channels[i], normalisation=norm, training=training) net = layers.conv2D(net, 'z%d_pre_3' % i, num_filters=num_channels[i], normalisation=norm, training=training) pre_z[i] = net # Generate z's for i in reversed(range(latent_levels)): spatial_zdim = [ d // 2**(i + resolution_levels - latent_levels) for d in spatial_xdim ] spatial_cov_dim = spatial_zdim[0] * spatial_zdim[1] if i == latent_levels - 1: mu[i] = layers.conv2D(pre_z[i + resolution_levels - latent_levels], 'z%d_mu' % i, num_filters=zdim_0, activation=tf.identity) if full_cov_list[i] == True: l = layers.dense_layer( pre_z[i + resolution_levels - latent_levels], 'z%d_sigma' % i, hidden_units=zdim_0 * spatial_cov_dim * (spatial_cov_dim + 1) // 2, activation=tf.identity) l = tf.reshape(l, [ -1, zdim_0, spatial_cov_dim * (spatial_cov_dim + 1) // 2 ]) Lp = tf.contrib.distributions.fill_triangular(l) L = tf.linalg.set_diag( Lp, tf.nn.softplus(tf.linalg.diag_part(Lp)) ) # Cholesky factors must have positive diagonal sigma[i] = L eps = tf.random_normal(tf.shape(mu[i])) eps = tf.transpose(eps, perm=[0, 3, 1, 2]) bs = tf.shape(x)[0] eps = tf.reshape(eps, tf.stack([bs, zdim_0, -1, 1])) eps_tmp = tf.matmul(sigma[i], eps) eps_tmp = tf.transpose(eps_tmp, perm=[0, 2, 3, 1]) eps_tmp = tf.reshape( eps_tmp, [bs, spatial_zdim[0], spatial_zdim[1], zdim_0]) z[i] = mu[i] + eps_tmp else: sigma[i] = layers.conv2D(pre_z[i + resolution_levels - latent_levels], 'z%d_sigma' % i, num_filters=zdim_0, activation=tf.nn.softplus, kernel_size=(1, 1)) z[i] = mu[i] + sigma[i] * tf.random_normal( tf.shape(mu[i]), 0, 1, dtype=tf.float32) else: for j in reversed(range(0, i + 1)): z_below_ups = layers.bilinear_upsample2D( z_ups_mat[j + 1][i + 1], factor=2, name='ups') z_below_ups = layers.conv2D(z_below_ups, name='z%d_ups_to_%d_c_1' % ((i + 1), (j + 1)), num_filters=zdim_0 * n0, normalisation=norm, training=training) z_below_ups = layers.conv2D(z_below_ups, name='z%d_ups_to_%d_c_2' % ((i + 1), (j + 1)), num_filters=zdim_0 * n0, normalisation=norm, training=training) z_ups_mat[j][i + 1] = z_below_ups if full_latent_dependencies: z_input = tf.concat( [pre_z[i + resolution_levels - latent_levels]] + z_ups_mat[i][(i + 1):latent_levels], axis=3, name='concat_%d' % i) else: z_input = tf.concat([ pre_z[i + resolution_levels - latent_levels], z_ups_mat[i][i + 1] ], axis=3, name='concat_%d' % i) z_input = layers.conv2D(z_input, 'z%d_input_1' % i, num_filters=num_channels[i], normalisation=norm, training=training) z_input = layers.conv2D(z_input, 'z%d_input_2' % i, num_filters=num_channels[i], normalisation=norm, training=training) mu[i] = layers.conv2D(z_input, 'z%d_mu' % i, num_filters=zdim_0, activation=tf.identity, kernel_size=(1, 1)) if full_cov_list[i] == True: l = layers.dense_layer(z_input, 'z%d_sigma' % i, hidden_units=zdim_0 * spatial_cov_dim * (spatial_cov_dim + 1) // 2, activation=tf.identity) l = tf.reshape(l, [ -1, zdim_0, spatial_cov_dim * (spatial_cov_dim + 1) // 2 ]) Lp = tf.contrib.distributions.fill_triangular(l) L = tf.linalg.set_diag( Lp, tf.nn.softplus(tf.linalg.diag_part(Lp))) sigma[i] = L eps = tf.random_normal(tf.shape(mu[i])) eps = tf.transpose(eps, perm=[0, 3, 1, 2]) bs = tf.shape(x)[0] eps = tf.reshape(eps, tf.stack([bs, zdim_0, -1, 1])) eps_tmp = tf.matmul(sigma[i], eps) eps_tmp = tf.transpose(eps_tmp, perm=[0, 2, 3, 1]) eps_tmp = tf.reshape( eps_tmp, [bs, spatial_zdim[0], spatial_zdim[1], zdim_0]) z[i] = mu[i] + eps_tmp else: sigma[i] = layers.conv2D(z_input, 'z%d_sigma' % i, num_filters=zdim_0, activation=tf.nn.softplus, kernel_size=(1, 1)) z[i] = mu[i] + sigma[i] * tf.random_normal( tf.shape(mu[i]), 0, 1, dtype=tf.float32) z_ups_mat[i][i] = z[i] return z, mu, sigma
def proposed(z_list, training, image_size, n_classes, scope_reuse=False, norm=tfnorm.batch_norm, rank=10, diagonal=False, **kwargs): x = kwargs.get('x') resolution_levels = kwargs.get('resolution_levels', 7) n0 = kwargs.get('n0', 32) num_channels = [n0, 2 * n0, 4 * n0, 6 * n0, 6 * n0, 6 * n0, 6 * n0] conv_unit = layers.conv2D deconv_unit = lambda inp: layers.bilinear_upsample2D(inp, 'upsample', 2) with tf.variable_scope('likelihood') as scope: if scope_reuse: scope.reuse_variables() add_bias = False if norm == tfnorm.batch_norm else True enc = [] with tf.variable_scope('encoder'): for ii in range(resolution_levels): enc.append([]) # In first layer set input to x rather than max pooling if ii == 0: enc[ii].append(x) else: enc[ii].append(layers.averagepool2D(enc[ii - 1][-1])) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_1' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_2' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_3' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) dec = [] with tf.variable_scope('decoder'): for jj in range(resolution_levels - 1): ii = resolution_levels - jj - 1 # used to index the encoder again dec.append([]) if jj == 0: next_inp = enc[ii][-1] else: next_inp = dec[jj - 1][-1] dec[jj].append(deconv_unit(next_inp)) # skip connection dec[jj].append( layers.crop_and_concat([dec[jj][-1], enc[ii - 1][-1]], axis=3)) dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_1' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias) ) # projection True to make it work with res units. dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_2' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_3' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) net = dec[-1][-1] recomb = conv_unit(net, 'recomb_0', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) recomb = conv_unit(recomb, 'recomb_1', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) recomb = conv_unit(recomb, 'recomb_2', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) epsilon = 1e-5 mean = layers.conv2D(recomb, 'mean', num_filters=n_classes, kernel_size=(1, 1), activation=tf.identity) log_cov_diag = layers.conv2D(recomb, 'diag', num_filters=n_classes, kernel_size=(1, 1), activation=tf.identity) cov_factor = layers.conv2D(recomb, 'factor', num_filters=n_classes * rank, kernel_size=(1, 1), activation=tf.identity) shape = image_size[:-1] + (n_classes, ) flat_size = np.prod(shape) mean = tf.reshape(mean, [-1, flat_size]) cov_diag = tf.exp(tf.reshape(log_cov_diag, [-1, flat_size])) + epsilon cov_factor = tf.reshape(cov_factor, [-1, flat_size, rank]) if diagonal: dist = DiagonalMultivariateNormal(loc=mean, cov_diag=cov_diag) else: dist = LowRankMultivariateNormal(loc=mean, cov_diag=cov_diag, cov_factor=cov_factor) s = dist.rsample((1, )) s = tf.reshape(s, (-1, ) + shape) return [[s], dist]
def prob_unet2D(z_list, training, image_size, n_classes, scope_reuse=False, norm=tfnorm.batch_norm, **kwargs): x = kwargs.get('x') z = z_list[0] resolution_levels = kwargs.get('resolution_levels', 7) n0 = kwargs.get('n0', 32) num_channels = [n0, 2 * n0, 4 * n0, 6 * n0, 6 * n0, 6 * n0, 6 * n0] conv_unit = layers.conv2D deconv_unit = lambda inp: layers.bilinear_upsample2D(inp, 'upsample', 2) bs = tf.shape(x)[0] zdim = z.get_shape().as_list()[-1] with tf.variable_scope('likelihood') as scope: if scope_reuse: scope.reuse_variables() add_bias = False if norm == tfnorm.batch_norm else True enc = [] with tf.variable_scope('encoder'): for ii in range(resolution_levels): enc.append([]) # In first layer set input to x rather than max pooling if ii == 0: enc[ii].append(x) else: enc[ii].append(layers.averagepool2D(enc[ii - 1][-1])) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_1' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_2' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_3' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) dec = [] with tf.variable_scope('decoder'): for jj in range(resolution_levels - 1): ii = resolution_levels - jj - 1 # used to index the encoder again dec.append([]) if jj == 0: next_inp = enc[ii][-1] else: next_inp = dec[jj - 1][-1] dec[jj].append(deconv_unit(next_inp)) # skip connection dec[jj].append( layers.crop_and_concat([dec[jj][-1], enc[ii - 1][-1]], axis=3)) dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_1' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias) ) # projection True to make it work with res units. dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_2' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_3' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) z_t = tf.reshape(z, tf.stack((bs, 1, 1, zdim))) broadcast_z = tf.tile(z_t, (1, image_size[0], image_size[1], 1)) net = tf.concat([dec[-1][-1], broadcast_z], axis=-1) recomb = conv_unit(net, 'recomb_0', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) recomb = conv_unit(recomb, 'recomb_1', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) recomb = conv_unit(recomb, 'recomb_2', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) s = [ layers.conv2D(recomb, 'prediction', num_filters=n_classes, kernel_size=(1, 1), activation=tf.identity) ] return s
def prob_unet2D_arch( x, training, nlabels, n0=32, resolution_levels=7, norm=tfnorm.batch_norm, conv_unit=layers.conv2D, deconv_unit=lambda inp: layers.bilinear_upsample2D(inp, 'upsample', 2), scope_reuse=False, return_net=False): num_channels = [n0, 2 * n0, 4 * n0, 6 * n0, 6 * n0, 6 * n0, 6 * n0] with tf.variable_scope('likelihood') as scope: if scope_reuse: scope.reuse_variables() add_bias = False if norm == tfnorm.batch_norm else True enc = [] with tf.variable_scope('encoder'): for ii in range(resolution_levels): enc.append([]) # In first layer set input to x rather than max pooling if ii == 0: enc[ii].append(x) else: enc[ii].append(layers.averagepool2D(enc[ii - 1][-1])) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_1' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_2' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) enc[ii].append( conv_unit(enc[ii][-1], 'conv_%d_3' % ii, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) dec = [] with tf.variable_scope('decoder'): for jj in range(resolution_levels - 1): ii = resolution_levels - jj - 1 # used to index the encoder again dec.append([]) if jj == 0: next_inp = enc[ii][-1] else: next_inp = dec[jj - 1][-1] dec[jj].append(deconv_unit(next_inp)) # skip connection dec[jj].append( layers.crop_and_concat([dec[jj][-1], enc[ii - 1][-1]], axis=3)) dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_1' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias) ) # projection True to make it work with res units. dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_2' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) dec[jj].append( conv_unit(dec[jj][-1], 'conv_%d_3' % jj, num_filters=num_channels[ii], training=training, normalisation=norm, add_bias=add_bias)) recomb = conv_unit(dec[-1][-1], 'recomb_0', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) recomb = conv_unit(recomb, 'recomb_1', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) recomb = conv_unit(recomb, 'recomb_2', num_filters=num_channels[0], kernel_size=(1, 1), training=training, normalisation=norm, add_bias=add_bias) s = layers.conv2D(recomb, 'prediction', num_filters=nlabels, kernel_size=(1, 1), activation=tf.identity) return s
def phiseg(x, s_oh, zdim_0, training, scope_reuse=False, norm=tfnorm.batch_norm, **kwargs): n0 = kwargs.get('n0', 32) num_channels = [n0, 2 * n0, 4 * n0, 6 * n0, 6 * n0, 6 * n0, 6 * n0] with tf.variable_scope('posterior') as scope: if scope_reuse: scope.reuse_variables() full_cov_list = kwargs.get('full_cov_list', None) n0 = kwargs.get('n0', 32) latent_levels = kwargs.get('latent_levels', 5) resolution_levels = kwargs.get('resolution_levels', 7) spatial_xdim = x.get_shape().as_list()[1:3] pre_z = [None] * resolution_levels mu = [None] * latent_levels sigma = [None] * latent_levels z = [None] * latent_levels z_ups_mat = [] for i in range(latent_levels): z_ups_mat.append( [None] * latent_levels) # encoding [original resolution][upsampled to] # Generate pre_z's for i in range(resolution_levels): if i == 0: net = tf.concat([x, s_oh - 0.5], axis=-1) else: net = layers.averagepool2D(pre_z[i - 1]) net = layers.conv2D(net, 'z%d_pre_1' % i, num_filters=num_channels[i], normalisation=norm, training=training) net = layers.conv2D(net, 'z%d_pre_2' % i, num_filters=num_channels[i], normalisation=norm, training=training) net = layers.conv2D(net, 'z%d_pre_3' % i, num_filters=num_channels[i], normalisation=norm, training=training) pre_z[i] = net # Generate z's for i in reversed(range(latent_levels)): spatial_zdim = [ d // 2**(i + resolution_levels - latent_levels) for d in spatial_xdim ] spatial_cov_dim = spatial_zdim[0] * spatial_zdim[1] if i == latent_levels - 1: mu[i] = layers.conv2D(pre_z[i + resolution_levels - latent_levels], 'z%d_mu' % i, num_filters=zdim_0, activation=tf.identity) sigma[i] = layers.conv2D(pre_z[i + resolution_levels - latent_levels], 'z%d_sigma' % i, num_filters=zdim_0, activation=tf.nn.softplus, kernel_size=(1, 1)) z[i] = mu[i] + sigma[i] * tf.random_normal( tf.shape(mu[i]), 0, 1, dtype=tf.float32) else: for j in reversed(range(0, i + 1)): z_below_ups = layers.bilinear_upsample2D( z_ups_mat[j + 1][i + 1], factor=2, name='ups') z_below_ups = layers.conv2D(z_below_ups, name='z%d_ups_to_%d_c_1' % ((i + 1), (j + 1)), num_filters=zdim_0 * n0, normalisation=norm, training=training) z_below_ups = layers.conv2D(z_below_ups, name='z%d_ups_to_%d_c_2' % ((i + 1), (j + 1)), num_filters=zdim_0 * n0, normalisation=norm, training=training) z_ups_mat[j][i + 1] = z_below_ups z_input = tf.concat([ pre_z[i + resolution_levels - latent_levels], z_ups_mat[i][i + 1] ], axis=3, name='concat_%d' % i) z_input = layers.conv2D(z_input, 'z%d_input_1' % i, num_filters=num_channels[i], normalisation=norm, training=training) z_input = layers.conv2D(z_input, 'z%d_input_2' % i, num_filters=num_channels[i], normalisation=norm, training=training) mu[i] = layers.conv2D(z_input, 'z%d_mu' % i, num_filters=zdim_0, activation=tf.identity, kernel_size=(1, 1)) sigma[i] = layers.conv2D(z_input, 'z%d_sigma' % i, num_filters=zdim_0, activation=tf.nn.softplus, kernel_size=(1, 1)) z[i] = mu[i] + sigma[i] * tf.random_normal( tf.shape(mu[i]), 0, 1, dtype=tf.float32) z_ups_mat[i][i] = z[i] return z, mu, sigma