Spaces:
Sleeping
Sleeping
File size: 6,214 Bytes
9016314 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
import numpy as np
import tensorflow as tf
initializer = tf.compat.v1.random_normal_initializer(stddev=0.001)
def dense_nn(x, dims, dim_out, norm=True, name='dense_nn'):
with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
for i, size in enumerate(dims):
x = tf.compat.v1.layers.dense(x, size, name=f'd{i}', kernel_initializer=initializer)
if norm:
x = tf.contrib.layers.layer_norm(x)
x = tf.nn.leaky_relu(x)
x = tf.compat.v1.layers.dense(x, dim_out, name='d_out', kernel_initializer=initializer)
return x
def cond_dense_nn(x, cond, dims, dim_out, norm=True, name='cond_dense_nn'):
with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
for i, size in enumerate(dims):
x = tf.compat.v1.layers.dense(x, size, name=f'd{i}', kernel_initializer=initializer)
if norm:
x = tf.contrib.layers.layer_norm(x)
x = tf.nn.leaky_relu(x)
c = tf.compat.v1.layers.dense(cond, size, name=f'c{i}', kernel_initializer=initializer)
x = tf.sigmoid(c) * x
x = tf.compat.v1.layers.dense(x, dim_out, name='d_out', kernel_initializer=initializer)
return x
def large_cond_dense_nn(x, cond, dims, dim_out, norm=True, name='cond_dense_nn'):
with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
for i, size in enumerate(dims):
x = tf.compat.v1.layers.dense(x, size, name=f'd{i}', kernel_initializer=initializer)
if norm:
x = tf.contrib.layers.layer_norm(x)
x = tf.nn.leaky_relu(x)
c = dense_nn(cond, [256,256], size, False, name=f'c{i}')
x = tf.sigmoid(c) * x
x = tf.compat.v1.layers.dense(x, dim_out, name='d_out')
return x
def res_block(input, dim, block_name):
x = tf.compat.v1.layers.dense(input, dim, name=f'{block_name}_1', kernel_initializer=initializer)
x = tf.nn.leaky_relu(x)
x = tf.compat.v1.layers.dense(input, dim, name=f'{block_name}_2', kernel_initializer=initializer)
x += input
x = tf.nn.leaky_relu(x)
return x
def cond_resnet(x, cond, dims, dim_out, norm=True, name='cond_resnet'):
with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
x = tf.compat.v1.layers.dense(x, dims[0], name='1', kernel_initializer=initializer)
for i, size in enumerate(dims):
x = res_block(x, size, block_name=f'res_block_{i}')
if norm:
x = tf.contrib.layers.layer_norm(x)
c = dense_nn(cond, [256,256], size, False, name=f'c{i}')
x = tf.sigmoid(c) * x
x = tf.compat.v1.layers.dense(x, dim_out, name='d_out', kernel_initializer=initializer)
return x
def convnet(x, dims, dim_out, name='convnet'):
with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
for i, d in enumerate(dims):
x = tf.compat.v1.layers.conv2d(x, d, 3, padding='same', name=f'c{i}_1')
x = tf.contrib.layers.instance_norm(x)
x = tf.nn.leaky_relu(x)
x = tf.compat.v1.layers.conv2d(x, d, 3, padding='same', name=f'c{i}_2')
x = tf.contrib.layers.instance_norm(x)
x = tf.nn.leaky_relu(x)
x = tf.compat.v1.layers.max_pooling2d(x, 2, 2)
x = tf.compat.v1.layers.flatten(x)
x = tf.compat.v1.layers.dense(x, dim_out, name='d1')
x = tf.contrib.layers.instance_norm(x)
x = tf.nn.leaky_relu(x)
x = tf.compat.v1.layers.dense(x, dim_out, name='d2')
return x
def peq_convnet(x, dims, dim_out, attention, name='peq_convnet'):
B,N,H,W,C = tf.shape(input=x)[0], tf.shape(input=x)[1], *x.get_shape().as_list()[2:]
with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
# downsample
x = tf.reshape(x, [-1,H,W,C])
for d in dims[:2]:
x = tf.compat.v1.layers.conv2d(x, d, 3, strides=(1,1), padding='same')
x = tf.contrib.layers.instance_norm(x)
x = tf.nn.leaky_relu(x)
x = tf.compat.v1.layers.conv2d(x, d, 3, strides=(2,2), padding='same')
x = tf.contrib.layers.instance_norm(x)
x = tf.nn.leaky_relu(x)
H, W, C = H//2, W//2, d
x = tf.reshape(x, [B,N,H,W,C])
# attention across set dimension
x = tf.reshape(tf.transpose(a=x, perm=[0,2,3,1,4]), [B*H*W,N,C])
rep = attention(x, x, x)
x += rep
x = tf.transpose(a=tf.reshape(x, [B,H,W,N,C]), perm=[0,3,1,2,4])
# downsample
x = tf.reshape(x, [-1,H,W,C])
for d in dims[2:]:
x = tf.compat.v1.layers.conv2d(x, d, 3, strides=(1,1), padding='same')
x = tf.contrib.layers.instance_norm(x)
x = tf.nn.leaky_relu(x)
x = tf.compat.v1.layers.conv2d(x, d, 3, strides=(2,2), padding='same')
x = tf.contrib.layers.instance_norm(x)
x = tf.nn.leaky_relu(x)
H, W, C = H//2, W//2, d
x = tf.compat.v1.layers.flatten(x)
x = tf.compat.v1.layers.dense(x, dim_out)
x = tf.contrib.layers.instance_norm(x)
x = tf.nn.leaky_relu(x)
x = tf.compat.v1.layers.dense(x, dim_out)
x = tf.reshape(x, [B,N,dim_out])
return x
def peq_resblock(x, dim, attention, name='peq_resnet'):
B,N,H,W,C = tf.shape(input=x)[0], tf.shape(input=x)[1], *x.get_shape().as_list()[2:]
with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
res = tf.reshape(tf.transpose(a=x, perm=[0,2,3,1,4]), [B*H*W,N,C])
res = attention(res, res, res)
res = tf.transpose(a=tf.reshape(res, [B,H,W,N,C]), perm=[0,3,1,2,4])
res = tf.reshape(res, [B*N,H,W,C])
res = tf.compat.v1.layers.conv2d(res, dim, 3, strides=(1,1), padding='same')
res = tf.contrib.layers.instance_norm(res)
res = tf.nn.leaky_relu(res)
res = tf.compat.v1.layers.conv2d(res, dim, 3, strides=(1,1), padding='same')
res = tf.contrib.layers.instance_norm(res)
res = tf.reshape(res, [B,N,H,W,C])
x += res
x = tf.nn.leaky_relu(x)
return x
|