import numpy as np import tensorflow as tf from absl import flags from libml import utils, ctaugment from libml.utils import EasyDict from third_party.auto_augment import augmentations, policies FLAGS = flags.FLAGS POOL = None POLICIES = EasyDict( cifar10=policies.cifar10_policies(), cifar10p=policies.cifar10_policies(), # color=policies.color_policies(), cifar10imb=policies.cifar10_policies(), cifar100=policies.cifar10_policies(), svhn=policies.svhn_policies(), svhnp=policies.svhn_policies(), svhnp_noextra=policies.svhn_policies(), svhn_noextra=policies.svhn_policies()) RANDOM_POLICY_OPS = ('Identity', 'AutoContrast', 'Equalize', 'Rotate', 'Solarize', 'Color', 'Contrast', 'Brightness', 'Sharpness', 'ShearX', 'TranslateX', 'TranslateY', 'Posterize', 'ShearY') AUGMENT_ENUM = 'd x m aa aac aacc ra rac'.split() + [ 'r%d_%d_%d' % (nops, mag, cutout) for nops, mag, cutout in itertools.product(range(1, 5), range(1, 16), range(0, 100, 25)) ] + ['rac%d' % (mag) for mag in range(1, 10)]
def classifier(self, x, scales, filters, repeat, training, getter=None, dropout=0, **kwargs): del kwargs leaky_relu = functools.partial(tf.nn.leaky_relu, alpha=0.1) bn_args = dict(training=training, momentum=0.999) def conv_args(k, f): return dict(padding='same', kernel_initializer=tf.random_normal_initializer( stddev=tf.rsqrt(0.5 * k * k * f))) def residual(x0, filters, stride=1, activate_before_residual=False): x = leaky_relu(tf.layers.batch_normalization(x0, **bn_args)) if activate_before_residual: x0 = x x = tf.layers.conv2d(x, filters, 3, strides=stride, **conv_args(3, filters)) x = leaky_relu(tf.layers.batch_normalization(x, **bn_args)) x = tf.layers.conv2d(x, filters, 3, **conv_args(3, filters)) if x0.get_shape()[3] != filters: x0 = tf.layers.conv2d(x0, filters, 1, strides=stride, **conv_args(1, filters)) return x0 + x with tf.variable_scope('classify', reuse=tf.AUTO_REUSE, custom_getter=getter): y = tf.layers.conv2d((x - self.dataset.mean) / self.dataset.std, 16, 3, **conv_args(3, 16)) for scale in range(scales): y = residual(y, filters << scale, stride=2 if scale else 1, activate_before_residual=scale == 0) for i in range(repeat - 1): y = residual(y, filters << scale) y = leaky_relu(tf.layers.batch_normalization(y, **bn_args)) y = embeds = tf.reduce_mean(y, [1, 2]) if dropout and training: y = tf.nn.dropout(y, 1 - dropout) logits = tf.layers.dense( y, self.nclass, kernel_initializer=tf.glorot_normal_initializer()) return EasyDict(logits=logits, embeds=embeds)
import itertools import multiprocessing import random import numpy as np import tensorflow as tf from absl import flags from libml import utils, ctaugment from libml.utils import EasyDict from third_party.auto_augment import augmentations, policies FLAGS = flags.FLAGS POOL = None POLICIES = EasyDict(cifar10=policies.cifar10_policies(), cifar100=policies.cifar10_policies(), svhn=policies.svhn_policies(), svhn_noextra=policies.svhn_policies()) RANDOM_POLICY_OPS = ('Identity', 'AutoContrast', 'Equalize', 'Rotate', 'Solarize', 'Color', 'Contrast', 'Brightness', 'Sharpness', 'ShearX', 'TranslateX', 'TranslateY', 'Posterize', 'ShearY') AUGMENT_ENUM = 'd x m aa aac ra rac'.split() + [ 'r%d_%d_%d' % (nops, mag, cutout) for nops, mag, cutout in itertools.product(range(1, 5), range(1, 16), range(0, 100, 25)) ] + ['rac%d' % (mag) for mag in range(1, 10)] flags.DEFINE_integer('K', 1, 'Number of strong augmentation for unlabeled data.') flags.DEFINE_enum(
def model(self, batch, lr, wd, beta, w_kl, w_match, w_rot, K, use_xe, warmup_kimg=1024, T=0.5, mixmode='xxy.yxy', dbuf=128, ema=0.999, **kwargs): hwc = [self.dataset.height, self.dataset.width, self.dataset.colors] xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt') # For training x_in = tf.placeholder(tf.float32, [None] + hwc, 'x') y_in = tf.placeholder(tf.float32, [batch, K + 1] + hwc, 'y') l_in = tf.placeholder(tf.int32, [batch], 'labels') w_match *= tf.clip_by_value( tf.cast(self.step, tf.float32) / (warmup_kimg << 10), 0, 1) lrate = tf.clip_by_value( tf.to_float(self.step) / (FLAGS.train_kimg << 10), 0, 1) lr *= tf.cos(lrate * (7 * np.pi) / (2 * 8)) tf.summary.scalar('monitors/lr', lr) augment = layers.MixMode(mixmode) gpu = utils.get_gpu() def classifier_to_gpu(x, **kw): with tf.device(next(gpu)): return self.classifier(x, **kw, **kwargs).logits def random_rotate(x): b4 = batch // 4 x, xt = x[:2 * b4], tf.transpose(x[2 * b4:], [0, 2, 1, 3]) l = np.zeros(b4, np.int32) l = tf.constant(np.concatenate([l, l + 1, l + 2, l + 3], axis=0)) return tf.concat( [x[:b4], x[b4:, ::-1, ::-1], xt[:b4, ::-1], xt[b4:, :, ::-1]], axis=0), l # Moving average of the current estimated label distribution p_model = layers.PMovingAverage('p_model', self.nclass, dbuf) p_target = layers.PMovingAverage( 'p_target', self.nclass, dbuf) # Rectified distribution (only for plotting) # Known (or inferred) true unlabeled distribution p_data = layers.PData(self.dataset) if w_rot > 0: rot_y, rot_l = random_rotate(y_in[:, 1]) with tf.device(next(gpu)): rot_logits = self.classifier_rot( self.classifier(rot_y, training=True, **kwargs).embeds) loss_rot = tf.nn.softmax_cross_entropy_with_logits_v2( labels=tf.one_hot(rot_l, 4), logits=rot_logits) loss_rot = tf.reduce_mean(loss_rot) tf.summary.scalar('losses/rot', loss_rot) else: loss_rot = 0 if kwargs['redux'] == '1st' and w_kl <= 0: logits_y = [classifier_to_gpu(y_in[:, 0], training=True)] * (K + 1) elif kwargs['redux'] == '1st': logits_y = [ classifier_to_gpu(y_in[:, i], training=True) for i in range(2) ] logits_y += logits_y[:1] * (K - 1) else: logits_y = [ classifier_to_gpu(y_in[:, i], training=True) for i in range(K + 1) ] guess = self.guess_label(logits_y, p_data(), p_model(), T=T, **kwargs) ly = tf.stop_gradient(guess.p_target) if w_kl > 0: w_kl *= tf.clip_by_value( tf.cast(self.step, tf.float32) / (warmup_kimg << 10), 0, 1) loss_kl = tf.nn.softmax_cross_entropy_with_logits_v2( labels=ly[:batch], logits=logits_y[1]) loss_kl = tf.reduce_mean(loss_kl) tf.summary.scalar('losses/kl', loss_kl) else: loss_kl = 0 del logits_y lx = tf.one_hot(l_in, self.nclass) xy, labels_xy = augment([xt_in] + [y_in[:, i] for i in range(K + 1)], [lx] + tf.split(ly, K + 1), [beta, beta]) x, y = xy[0], xy[1:] labels_x, labels_y = labels_xy[0], tf.concat(labels_xy[1:], 0) del xy, labels_xy batches = layers.interleave([x] + y, batch) logits = [classifier_to_gpu(yi, training=True) for yi in batches[:-1]] skip_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) logits.append(classifier_to_gpu(batches[-1], training=True)) post_ops = [ v for v in tf.get_collection(tf.GraphKeys.UPDATE_OPS) if v not in skip_ops ] logits = layers.interleave(logits, batch) logits_x = logits[0] logits_y = tf.concat(logits[1:], 0) del batches, logits loss_xe = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels_x, logits=logits_x) loss_xe = tf.reduce_mean(loss_xe) if use_xe: loss_xeu = tf.nn.softmax_cross_entropy_with_logits_v2( labels=labels_y, logits=logits_y) else: loss_xeu = tf.square(labels_y - tf.nn.softmax(logits_y)) loss_xeu = tf.reduce_mean(loss_xeu) tf.summary.scalar('losses/xe', loss_xe) tf.summary.scalar('losses/%s' % ('xeu' if use_xe else 'l2u'), loss_xeu) self.distribution_summary(p_data(), p_model(), p_target()) # L2 regularization loss_wd = sum( tf.nn.l2_loss(v) for v in utils.model_vars('classify') if 'kernel' in v.name) tf.summary.scalar('losses/wd', loss_wd) ema = tf.train.ExponentialMovingAverage(decay=ema) ema_op = ema.apply(utils.model_vars()) ema_getter = functools.partial(utils.getter_ema, ema) post_ops.extend([ ema_op, p_model.update(guess.p_model), p_target.update(guess.p_target) ]) if p_data.has_update: post_ops.append(p_data.update(lx)) train_op = tf.train.MomentumOptimizer( lr, 0.9, use_nesterov=True).minimize( loss_xe + w_kl * loss_kl + w_match * loss_xeu + w_rot * loss_rot + wd * loss_wd, colocate_gradients_with_ops=True) with tf.control_dependencies([train_op]): train_op = tf.group(*post_ops) return EasyDict( xt=xt_in, x=x_in, y=y_in, label=l_in, train_op=train_op, classify_op=tf.nn.softmax( classifier_to_gpu(x_in, getter=ema_getter, training=False)), classify_raw=tf.nn.softmax(classifier_to_gpu( x_in, training=False))) # No EMA, for debugging.
def classifier(self, x, scales, filters, repeat, training, getter=None, dropout=0, **kwargs): del kwargs bn_args = dict(training=training, momentum=0.999) def conv_args(k, f): return dict(padding='same', use_bias=False, kernel_initializer=tf.random_normal_initializer( stddev=tf.rsqrt(0.5 * k * k * f))) def residual(x0, filters, stride=1): def branch(): x = tf.nn.relu(x0) x = tf.layers.conv2d(x, filters, 3, strides=stride, **conv_args(3, filters)) x = tf.nn.relu(tf.layers.batch_normalization(x, **bn_args)) x = tf.layers.conv2d(x, filters, 3, **conv_args(3, filters)) x = tf.layers.batch_normalization(x, **bn_args) return x x = layers.shakeshake(branch(), branch(), training) if stride == 2: x1 = tf.layers.conv2d(tf.nn.relu(x0[:, ::2, ::2]), filters >> 1, 1, **conv_args(1, filters >> 1)) x2 = tf.layers.conv2d(tf.nn.relu(x0[:, 1::2, 1::2]), filters >> 1, 1, **conv_args(1, filters >> 1)) x0 = tf.concat([x1, x2], axis=3) x0 = tf.layers.batch_normalization(x0, **bn_args) elif x0.get_shape()[3] != filters: x0 = tf.layers.conv2d(x0, filters, 1, **conv_args(1, filters)) x0 = tf.layers.batch_normalization(x0, **bn_args) return x0 + x with tf.variable_scope('classify', reuse=tf.AUTO_REUSE, custom_getter=getter): y = tf.layers.conv2d((x - self.dataset.mean) / self.dataset.std, 16, 3, **conv_args(3, 16)) for scale, i in itertools.product(range(scales), range(repeat)): with tf.variable_scope('layer%d.%d' % (scale + 1, i)): if i == 0: y = residual(y, filters << scale, stride=2 if scale else 1) else: y = residual(y, filters << scale) y = embeds = tf.reduce_mean(y, [1, 2]) if dropout and training: y = tf.nn.dropout(y, 1 - dropout) logits = tf.layers.dense( y, self.nclass, kernel_initializer=tf.glorot_normal_initializer()) return EasyDict(logits=logits, embeds=embeds)
def model(self, batch, lr, wd, ema, warmup_pos, consistency_weight, beta, **kwargs): hwc = [self.dataset.height, self.dataset.width, self.dataset.colors] xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt') # For training x_in = tf.placeholder(tf.float32, [None] + hwc, 'x') y_in = tf.placeholder(tf.float32, [batch, 2] + hwc, 'y') l_in = tf.placeholder(tf.int32, [batch], 'labels') l = tf.one_hot(l_in, self.nclass) wd *= lr warmup = tf.clip_by_value( tf.to_float(self.step) / (warmup_pos * (FLAGS.train_kimg << 10)), 0, 1) y = tf.reshape(tf.transpose(y_in, [1, 0, 2, 3, 4]), [-1] + hwc) y_1, y_2 = tf.split(y, 2) mix = tf.distributions.Beta(beta, beta).sample([tf.shape(x_in)[0], 1, 1, 1]) mix = tf.maximum(mix, 1 - mix) classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits logits_x = classifier(xt_in, training=True) post_ops = tf.get_collection( tf.GraphKeys.UPDATE_OPS ) # Take only first call to update batch norm. ema = tf.train.ExponentialMovingAverage(decay=ema) ema_op = ema.apply(utils.model_vars()) ema_getter = functools.partial(utils.getter_ema, ema) logits_teacher = classifier(y_1, training=True, getter=ema_getter) labels_teacher = tf.stop_gradient(tf.nn.softmax(logits_teacher)) labels_teacher = labels_teacher * mix[:, :, 0, 0] + labels_teacher[::-1] * ( 1 - mix[:, :, 0, 0]) logits_student = classifier(y_1 * mix + y_1[::-1] * (1 - mix), training=True) loss_mt = tf.reduce_mean( (labels_teacher - tf.nn.softmax(logits_student))**2, -1) loss_mt = tf.reduce_mean(loss_mt) loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits_x) loss = tf.reduce_mean(loss) tf.summary.scalar('losses/xe', loss) tf.summary.scalar('losses/mt', loss_mt) post_ops.append(ema_op) post_ops.extend([ tf.assign(v, v * (1 - wd)) for v in utils.model_vars('classify') if 'kernel' in v.name ]) train_op = tf.train.AdamOptimizer(lr).minimize( loss + loss_mt * warmup * consistency_weight, colocate_gradients_with_ops=True) with tf.control_dependencies([train_op]): train_op = tf.group(*post_ops) return EasyDict( xt=xt_in, x=x_in, y=y_in, label=l_in, train_op=train_op, classify_raw=tf.nn.softmax(classifier( x_in, training=False)), # No EMA, for debugging. classify_op=tf.nn.softmax( classifier(x_in, getter=ema_getter, training=False)))
def model(self, batch, lr, wd, ema, beta, w_match, warmup_kimg=1024, nu=2, mixmode='xxy.yxy', dbuf=128, **kwargs): hwc = [self.dataset.height, self.dataset.width, self.dataset.colors] xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt') # For training x_in = tf.placeholder(tf.float32, [None] + hwc, 'x') y_in = tf.placeholder(tf.float32, [batch, nu] + hwc, 'y') l_in = tf.placeholder(tf.int32, [batch], 'labels') w_match *= tf.clip_by_value( tf.cast(self.step, tf.float32) / (warmup_kimg << 10), 0, 1) lrate = tf.clip_by_value( tf.to_float(self.step) / (FLAGS.train_kimg << 10), 0, 1) lr *= tf.cos(lrate * (7 * np.pi) / (2 * 8)) tf.summary.scalar('monitors/lr', lr) augment = MixMode(mixmode) classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits # Moving average of the current estimated label distribution p_model = layers.PMovingAverage('p_model', self.nclass, dbuf) p_target = layers.PMovingAverage( 'p_target', self.nclass, dbuf) # Rectified distribution (only for plotting) # Known (or inferred) true unlabeled distribution p_data = layers.PData(self.dataset) y = tf.reshape(tf.transpose(y_in, [1, 0, 2, 3, 4]), [-1] + hwc) guess = self.guess_label(tf.split(y, nu), classifier, T=0.5, **kwargs) ly = tf.stop_gradient(guess.p_target) lx = tf.one_hot(l_in, self.nclass) xy, labels_xy = augment([xt_in] + tf.split(y, nu), [lx] + [ly] * nu, [beta, beta]) x, y = xy[0], xy[1:] labels_x, labels_y = labels_xy[0], tf.concat(labels_xy[1:], 0) del xy, labels_xy batches = layers.interleave([x] + y, batch) skip_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) logits = [classifier(batches[0], training=True)] post_ops = [ v for v in tf.get_collection(tf.GraphKeys.UPDATE_OPS) if v not in skip_ops ] for batchi in batches[1:]: logits.append(classifier(batchi, training=True)) logits = layers.interleave(logits, batch) logits_x = logits[0] logits_y = tf.concat(logits[1:], 0) loss_xe = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels_x, logits=logits_x) loss_xe = tf.reduce_mean(loss_xe) loss_l2u = tf.square(labels_y - tf.nn.softmax(logits_y)) loss_l2u = tf.reduce_mean(loss_l2u) tf.summary.scalar('losses/xe', loss_xe) tf.summary.scalar('losses/l2u', loss_l2u) self.distribution_summary(p_data(), p_model(), p_target()) # L2 regularization loss_wd = sum( tf.nn.l2_loss(v) for v in utils.model_vars('classify') if 'kernel' in v.name) tf.summary.scalar('losses/wd', loss_wd) ema = tf.train.ExponentialMovingAverage(decay=ema) ema_op = ema.apply(utils.model_vars()) ema_getter = functools.partial(utils.getter_ema, ema) post_ops.append(ema_op) train_op = tf.train.MomentumOptimizer( lr, 0.9, use_nesterov=True).minimize( loss_xe + w_match * loss_l2u + wd * loss_wd, colocate_gradients_with_ops=True) with tf.control_dependencies([train_op]): train_op = tf.group(*post_ops) return EasyDict( xt=xt_in, x=x_in, y=y_in, label=l_in, train_op=train_op, classify_raw=tf.nn.softmax(classifier( x_in, training=False)), # No EMA, for debugging. classify_op=tf.nn.softmax( classifier(x_in, getter=ema_getter, training=False)))
def model(self, batch, lr, wd, ema, warmup_pos, consistency_weight, **kwargs): hwc = [self.dataset.height, self.dataset.width, self.dataset.colors] xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt') # For training x_in = tf.placeholder(tf.float32, [None] + hwc, 'x') y_in = tf.placeholder(tf.float32, [batch, 2] + hwc, 'y') l_in = tf.placeholder(tf.int32, [batch], 'labels') l = tf.one_hot(l_in, self.nclass) warmup = tf.clip_by_value( tf.to_float(self.step) / (warmup_pos * (FLAGS.train_kimg << 10)), 0, 1) lrate = tf.clip_by_value( tf.to_float(self.step) / (FLAGS.train_kimg << 10), 0, 1) lr *= tf.cos(lrate * (7 * np.pi) / (2 * 8)) tf.summary.scalar('monitors/lr', lr) classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits logits_x = classifier(xt_in, training=True) post_ops = tf.get_collection( tf.GraphKeys.UPDATE_OPS ) # Take only first call to update batch norm. y = tf.reshape(tf.transpose(y_in, [1, 0, 2, 3, 4]), [-1] + hwc) y_1, y_2 = tf.split(y, 2) ema = tf.train.ExponentialMovingAverage(decay=ema) ema_op = ema.apply(utils.model_vars()) ema_getter = functools.partial(utils.getter_ema, ema) logits_y = classifier(y_1, training=True, getter=ema_getter) logits_teacher = tf.stop_gradient(logits_y) logits_student = classifier(y_2, training=True) loss_mt = tf.reduce_mean( (tf.nn.softmax(logits_teacher) - tf.nn.softmax(logits_student))**2, -1) loss_mt = tf.reduce_mean(loss_mt) loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits_x) loss = tf.reduce_mean(loss) tf.summary.scalar('losses/xe', loss) tf.summary.scalar('losses/mt', loss_mt) # L2 regularization loss_wd = sum( tf.nn.l2_loss(v) for v in utils.model_vars('classify') if 'kernel' in v.name) tf.summary.scalar('losses/wd', loss_wd) post_ops.append(ema_op) train_op = tf.train.MomentumOptimizer( lr, 0.9, use_nesterov=True).minimize( loss + loss_mt * warmup * consistency_weight + wd * loss_wd, colocate_gradients_with_ops=True) with tf.control_dependencies([train_op]): train_op = tf.group(*post_ops) return EasyDict( xt=xt_in, x=x_in, y=y_in, label=l_in, train_op=train_op, classify_raw=tf.nn.softmax(classifier( x_in, training=False)), # No EMA, for debugging. classify_op=tf.nn.softmax( classifier(x_in, getter=ema_getter, training=False)))
def model(self, batch, lr, wd, ema, beta, w_match, warmup_kimg=1024, nu=2, mixmode='xxy.yxy', dbuf=128, **kwargs): # height, width, colors hwc = [self.dataset.height, self.dataset.width, self.dataset.colors] # labeled data [batch,32,32,3] xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt') # For training # labeled data [?,32,32,3] x_in = tf.placeholder(tf.float32, [None] + hwc, 'x') # unlabeled data [?,2,32,32,3], 每个未标记样本生成两个数据增强样本(nu=2) y_in = tf.placeholder(tf.float32, [batch, nu] + hwc, 'y') l_in = tf.placeholder(tf.int32, [batch], 'labels') # 使用weight decay调整权重,防止过拟合 wd *= lr # 在训练的前期逐步让w_match增长到最大值 w_match *= tf.clip_by_value( tf.cast(self.step, tf.float32) / (warmup_kimg << 10), 0, 1) augment = MixMode(mixmode) classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits # Moving average of the current estimated label distribution p_model = layers.PMovingAverage('p_model', self.nclass, dbuf) p_target = layers.PMovingAverage( 'p_target', self.nclass, dbuf) # Rectified distribution (only for plotting) # Known (or inferred) true unlabeled distribution p_data = layers.PData(self.dataset) # 将y_in([batch, nu, hwc]) 转化为[nu * batch, hwc] y = tf.reshape(tf.transpose(y_in, [1, 0, 2, 3, 4]), [-1] + hwc) guess = self.guess_label(tf.split(y, nu), classifier, T=0.5, **kwargs) # ly:unlabeled data的目标标签,当做真实标签使用 ly = tf.stop_gradient(guess.p_target) lx = tf.one_hot(l_in, self.nclass) xy, labels_xy = augment([xt_in] + tf.split(y, nu), [lx] + [ly] * nu, [beta, beta]) x, y = xy[0], xy[1:] labels_x, labels_y = labels_xy[0], tf.concat(labels_xy[1:], 0) del xy, labels_xy batches = layers.interleave([x] + y, batch) skip_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) logits = [classifier(batches[0], training=True)] post_ops = [ v for v in tf.get_collection(tf.GraphKeys.UPDATE_OPS) if v not in skip_ops ] for batchi in batches[1:]: logits.append(classifier(batchi, training=True)) logits = layers.interleave(logits, batch) logits_x = logits[0] logits_y = tf.concat(logits[1:], 0) loss_xe = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels_x, logits=logits_x) loss_xe = tf.reduce_mean(loss_xe) loss_l2u = tf.square(labels_y - tf.nn.softmax(logits_y)) loss_l2u = tf.reduce_mean(loss_l2u) tf.summary.scalar('losses/xe', loss_xe) tf.summary.scalar('losses/l2u', loss_l2u) self.distribution_summary(p_data(), p_model(), p_target()) ema = tf.train.ExponentialMovingAverage(decay=ema) ema_op = ema.apply(utils.model_vars()) ema_getter = functools.partial(utils.getter_ema, ema) post_ops.extend([ ema_op, p_model.update(guess.p_model), p_target.update(guess.p_target) ]) if p_data.has_update: post_ops.append(p_data.update(lx)) post_ops.extend([ tf.assign(v, v * (1 - wd)) for v in utils.model_vars('classify') if 'kernel' in v.name ]) train_op = tf.train.AdamOptimizer(lr).minimize( loss_xe + w_match * loss_l2u, colocate_gradients_with_ops=True) with tf.control_dependencies([train_op]): train_op = tf.group(*post_ops) return EasyDict( xt=xt_in, x=x_in, y=y_in, label=l_in, train_op=train_op, classify_raw=tf.nn.softmax(classifier( x_in, training=False)), # No EMA, for debugging. classify_op=tf.nn.softmax( classifier(x_in, getter=ema_getter, training=False)))
def model(self, batch, lr, wd, ema, warmup_pos, consistency_weight, threshold, **kwargs): hwc = [self.dataset.height, self.dataset.width, self.dataset.colors] xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt') # For training x_in = tf.placeholder(tf.float32, [None] + hwc, 'x') y_in = tf.placeholder(tf.float32, [batch] + hwc, 'y') l_in = tf.placeholder(tf.int32, [batch], 'labels') l = tf.one_hot(l_in, self.nclass) wd *= lr warmup = tf.clip_by_value( tf.to_float(self.step) / (warmup_pos * (FLAGS.train_kimg << 10)), 0, 1) classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits logits_x = classifier(xt_in, training=True) post_ops = tf.get_collection( tf.GraphKeys.UPDATE_OPS ) # Take only first call to update batch norm. logits_y = classifier(y_in, training=True) # Get the pseudo-label loss loss_pl = tf.nn.sparse_softmax_cross_entropy_with_logits( labels=tf.argmax(logits_y, axis=-1), logits=logits_y) # Masks denoting which data points have high-confidence predictions greater_than_thresh = tf.reduce_any( tf.greater(tf.nn.softmax(logits_y), threshold), axis=-1, keepdims=True, ) greater_than_thresh = tf.cast(greater_than_thresh, loss_pl.dtype) # Only enforce the loss when the model is confident loss_pl *= greater_than_thresh # Note that we also average over examples without confident outputs; # this is consistent with the realistic evaluation codebase loss_pl = tf.reduce_mean(loss_pl) loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits_x) loss = tf.reduce_mean(loss) tf.summary.scalar('losses/xe', loss) tf.summary.scalar('losses/pl', loss_pl) ema = tf.train.ExponentialMovingAverage(decay=ema) ema_op = ema.apply(utils.model_vars()) ema_getter = functools.partial(utils.getter_ema, ema) post_ops.append(ema_op) post_ops.extend([ tf.assign(v, v * (1 - wd)) for v in utils.model_vars('classify') if 'kernel' in v.name ]) train_op = tf.train.AdamOptimizer(lr).minimize( loss + loss_pl * warmup * consistency_weight, colocate_gradients_with_ops=True) with tf.control_dependencies([train_op]): train_op = tf.group(*post_ops) return EasyDict( xt=xt_in, x=x_in, y=y_in, label=l_in, train_op=train_op, classify_raw=tf.nn.softmax(classifier( x_in, training=False)), # No EMA, for debugging. classify_op=tf.nn.softmax( classifier(x_in, getter=ema_getter, training=False)))