def optimizer(self):
lr = tf.get_variable('learning_rate', initializer=0.1, trainable=False)
tf.summary.scalar('learning_rate-summary', lr)
opt = tf.train.MomentumOptimizer(lr, 0.9, use_nesterov=True)
if self.accum_grad != 1:
opt = AccumGradOptimizer(opt, self.accum_grad)
return opt
def optimizer(self):
lr = tf.get_variable('learning_rate', initializer=5e-4, trainable=False)
opt = tf.train.AdamOptimizer(lr, epsilon=1e-3)
return optimizer.apply_grad_processors(
opt, [
gradproc.ScaleGradient(('STN.*', 0.1)),
gradproc.SummaryGradient()])
def optimizer(self):
lr = tf.get_variable('learning_rate', initializer=0., trainable=False)
tf.summary.scalar('learning_rate-summary', lr)
# The learning rate in the config is set for 8 GPUs, and we use trainers with average=False.
lr = lr / 8.
opt = tf.train.MomentumOptimizer(lr, 0.9)
if cfg.TRAIN.NUM_GPUS < 8:
opt = optimizer.AccumGradOptimizer(opt, 8 // cfg.TRAIN.NUM_GPUS)
return opt
def DepthConv(x, out_channel, kernel_shape, padding='SAME', stride=1,
W_init=None, activation=tf.identity):
in_shape = x.get_shape().as_list()
in_channel = in_shape[1]
assert out_channel % in_channel == 0, (out_channel, in_channel)
channel_mult = out_channel // in_channel
if W_init is None:
W_init = tf.variance_scaling_initializer(2.0)
kernel_shape = [kernel_shape, kernel_shape]
filter_shape = kernel_shape + [in_channel, channel_mult]
W = tf.get_variable('W', filter_shape, initializer=W_init)
conv = tf.nn.depthwise_conv2d(x, W, [1, 1, stride, stride], padding=padding, data_format='NCHW')
return activation(conv, name='output')
def make_prediction(self, img):
img = tf.cast(img, tf.float32)
img = tf.image.rgb_to_grayscale(img)
k = tf.get_variable('filter',
dtype=tf.float32,
initializer=[[[[0.]], [[1.]], [[0.]]],
[[[1.]], [[-4.]], [[1.]]],
[[[0.]], [[1.]], [[0.]]]])
prediction_img = tf.nn.conv2d(img,
k,
strides=[1, 1, 1, 1],
padding='SAME')
return prediction_img
def GroupNorm(x, group, gamma_initializer=tf.constant_initializer(1.)):
"""
https://arxiv.org/abs/1803.08494
More code that reproduces the paper can be found at https://github.com/ppwwyyxx/GroupNorm-reproduce/.
"""
shape = x.get_shape().as_list()
ndims = len(shape)
assert ndims == 4, shape
chan = shape[1]
assert chan % group == 0, chan
group_size = chan // group
orig_shape = tf.shape(x)
h, w = orig_shape[2], orig_shape[3]
x = tf.reshape(x, tf.stack([-1, group, group_size, h, w]))
mean, var = tf.nn.moments(x, [2, 3, 4], keep_dims=True)
new_shape = [1, group, group_size, 1, 1]
beta = tf.get_variable('beta', [chan],
initializer=tf.constant_initializer())
beta = tf.reshape(beta, new_shape)
gamma = tf.get_variable('gamma', [chan], initializer=gamma_initializer)
gamma = tf.reshape(gamma, new_shape)
out = tf.nn.batch_normalization(x,
mean,
var,
beta,
gamma,
1e-5,
name='output')
return tf.reshape(out, orig_shape, name='output')
train_files.reset_state()
all_train_files = list(train_files)
all_train_files = all_train_files[:len(all_train_files) // args.batch *
args.batch] # truncate
num_train_images = len(all_train_files)
logger.info(
f"Creating graph for KNN of {num_train_images} training images ...")
local_train_files = [(idx, fname, label)
for idx, (fname, label) in enumerate(all_train_files)
if idx % hvd.size() == hvd.rank()]
image_input = tf.placeholder(tf.uint8, [None, 224, 224, 3], "image")
idx_input = tf.placeholder(tf.int64, [None], "image_idx")
feat_buffer = tf.get_variable("feature_buffer",
shape=[num_train_images, 128],
trainable=False)
net = ResNetModel(num_output=(2048, 128) if args.v2 else (128, ))
with TowerContext("", is_training=False):
feat = net.forward(image_input)
feat = tf.math.l2_normalize(feat, axis=1) # Nx128
all_feat = hvd.allgather(feat) # GN x 128
all_idx_input = hvd.allgather(idx_input) # GN
update_buffer = tf.scatter_update(feat_buffer, all_idx_input, all_feat)
dist = tf.matmul(feat, tf.transpose(feat_buffer)) # N x #DS
_, topk_indices = tf.math.top_k(dist, k=args.top_k) # Nxtopk
train_ds = build_dataflow(local_train_files)
config = get_default_sess_config()
def optimizer(self):
lr = tf.get_variable('learning_rate',
initializer=0.045,
trainable=False)
return tf.train.MomentumOptimizer(lr, 0.9)
def optimizer(self):
lr = tf.get_variable('learning_rate',
initializer=1e-2,
trainable=False)
tf.summary.scalar('lr', lr)
return tf.train.AdamOptimizer(lr, epsilon=1e-3)
def optimizer(self):
lr = tf.get_variable('learning_rate', initializer=0.0, trainable=False)
return tf.train.AdamOptimizer(lr)