def build_graph(self, image, label):
image = self.image_preprocess(image)
assert self.data_format in ['NCHW', 'NHWC']
if self.data_format == 'NCHW':
image = tf.transpose(image, [0, 3, 1, 2])
logits = self.get_logits(image)
tf.nn.softmax(logits, name='prob')
loss = ImageNetModel.compute_loss_and_error(
logits, label, label_smoothing=self.label_smoothing)
if self.weight_decay > 0:
wd_loss = regularize_cost(self.weight_decay_pattern,
l2_regularizer(self.weight_decay),
name='l2_regularize_loss')
add_moving_summary(loss, wd_loss)
total_cost = tf.add_n([loss, wd_loss], name='cost')
else:
total_cost = tf.identity(loss, name='cost')
add_moving_summary(total_cost)
if self.loss_scale != 1.:
logger.info("Scaling the total loss by {} ...".format(self.loss_scale))
return total_cost * self.loss_scale
else:
return total_cost
def channel_shuffle(l, group):
in_shape = l.get_shape().as_list()
in_channel = in_shape[1]
assert in_channel % group == 0, in_channel
l = tf.reshape(l, [-1, in_channel // group, group] + in_shape[-2:])
l = tf.transpose(l, [0, 2, 1, 3, 4])
l = tf.reshape(l, [-1, in_channel] + in_shape[-2:])
return l
def get_stn(image):
stn = (LinearWrap(image)
.AvgPooling('downsample', 2)
.Conv2D('conv0', 20, 5, padding='VALID')
.MaxPooling('pool0', 2)
.Conv2D('conv1', 20, 5, padding='VALID')
.FullyConnected('fc1', 32)
.FullyConnected('fct', 6, activation=tf.identity,
kernel_initializer=tf.constant_initializer(),
bias_initializer=tf.constant_initializer([1, 0, HALF_DIFF, 0, 1, HALF_DIFF]))())
# output 6 parameters for affine transformation
stn = tf.reshape(stn, [-1, 2, 3], name='affine') # bx2x3
stn = tf.reshape(tf.transpose(stn, [2, 0, 1]), [3, -1]) # 3 x (bx2)
coor = tf.reshape(tf.matmul(xys, stn),
[WARP_TARGET_SIZE, WARP_TARGET_SIZE, -1, 2])
coor = tf.transpose(coor, [2, 0, 1, 3], 'sampled_coords') # b h w 2
sampled = GridSample('warp', [image, coor], borderMode='constant')
return sampled
def visualize_conv1_weights(filters):
ctx = get_current_tower_context()
if not ctx.is_main_training_tower:
return
with tf.name_scope('visualize_conv1'):
filters = tf.reshape(filters, [11, 11, 3, 8, 12])
filters = tf.transpose(filters, [3, 0, 4, 1, 2]) # 8,11,12,11,3
filters = tf.reshape(filters, [1, 88, 132, 3])
tf.summary.image('visualize_conv1',
filters,
max_outputs=1,
collections=['AAA'])
def build_graph(self, image, label):
drop_rate = tf.constant(0.5 if self.training else 0.0)
if self.training:
tf.summary.image("train_image", image, 10)
if tf.test.is_gpu_available():
image = tf.transpose(image, [0, 3, 1, 2])
data_format = 'channels_first'
else:
data_format = 'channels_last'
image = image / 4.0 # just to make range smaller
with argscope(Conv2D, activation=BNReLU, use_bias=False, kernel_size=3), \
argscope([Conv2D, MaxPooling, BatchNorm], data_format=data_format):
logits = LinearWrap(image) \
.Conv2D('conv1.1', filters=64) \
.Conv2D('conv1.2', filters=64) \
.MaxPooling('pool1', 3, stride=2, padding='SAME') \
.Conv2D('conv2.1', filters=128) \
.Conv2D('conv2.2', filters=128) \
.MaxPooling('pool2', 3, stride=2, padding='SAME') \
.Conv2D('conv3.1', filters=128, padding='VALID') \
.Conv2D('conv3.2', filters=128, padding='VALID') \
.FullyConnected('fc0', 1024 + 512, activation=tf.nn.relu) \
.Dropout(rate=drop_rate) \
.FullyConnected('fc1', 512, activation=tf.nn.relu) \
.FullyConnected('linear', out_dim=self.cifar_classnum)()
cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
labels=label)
cost = tf.reduce_mean(cost, name='cross_entropy_loss')
correct = tf.cast(tf.nn.in_top_k(predictions=logits,
targets=label,
k=1),
tf.float32,
name='correct')
# monitor training error
add_moving_summary(tf.reduce_mean(correct, name='accuracy'))
# weight decay on all W of fc layers
wd_cost = regularize_cost('fc.*/W',
l2_regularizer(4e-4),
name='regularize_loss')
add_moving_summary(cost, wd_cost)
add_param_summary(('.*/W', ['histogram'])) # monitor W
return tf.add_n([cost, wd_cost], name='cost')
def visualize_conv_weights(filters, name):
"""Visualize use weights in convolution filters.
Args:
filters: tensor containing the weights [H,W,Cin,Cout]
name: label for tensorboard
Returns:
image of all weight
"""
with tf.name_scope('visualize_w_' + name):
filters = tf.transpose(
filters, (3, 2, 0, 1)) # [h, w, cin, cout] -> [cout, cin, h, w]
filters = tf.unstack(filters) # --> cout * [cin, h, w]
filters = tf.concat(filters, 1) # --> [cin, cout * h, w]
filters = tf.unstack(filters) # --> cin * [cout * h, w]
filters = tf.concat(filters, 1) # --> [cout * h, cin * w]
filters = tf.expand_dims(filters, 0)
filters = tf.expand_dims(filters, -1)
tf.summary.image('visualize_w_' + name, filters)
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()
config.gpu_options.visible_device_list = str(hvd.local_rank())
def evaluate(checkpoint_file):
result_file = get_checkpoint_path(
checkpoint_file) + f".knn{args.top_k}.txt"
if os.path.isfile(result_file):
logger.info(f"Skipping evaluation of {result_file}.")
return
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
def preprocess(self, image):
image = tf.expand_dims(image, 0)
image = image_preprocess(image, bgr=True)
return tf.transpose(image, [0, 3, 1, 2])