def main():
batch_size = config['batch_size']
experiment_dir = config['experiment_dir']
# setup experiment and checkpoint directories
checkpoint_dir = pth.join(experiment_dir, 'checkpoints')
if not pth.exists(experiment_dir):
os.makedirs(experiment_dir)
if not pth.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
trn_data_generator, vld_data = dataset.get_cifar10(batch_size)
train(trn_data_generator, vld_data)
def main(argv=None):
num_gpus = config['num_gpus']
batch_size = config['batch_size']
checkpoint_dir = config["checkpoint_dir"]
experiment_dir = config["experiment_dir"]
# setup experiment and checkpoint directories
if not pth.exists(experiment_dir):
os.makedirs(experiment_dir)
if not pth.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
train_data_generator, valset = dataset.get_cifar10(batch_size * num_gpus)
train(train_data_generator)
def main():
global best_acc
args.out = args.dataset + '@N_' + str(args.num_max) + '_r_'
if args.imb_ratio_l == args.imb_ratio_u:
args.out += str(args.imb_ratio_l) + '_' + args.semi_method
else:
args.out += str(args.imb_ratio_l) + '_' + str(
args.imb_ratio_u) + '_' + args.semi_method
if args.darp:
args.out += '_darp_alpha' + str(args.alpha) + '_iterT' + str(
args.iter_T)
if not os.path.isdir(args.out):
mkdir_p(args.out)
# Data
N_SAMPLES_PER_CLASS = make_imb_data(args.num_max, args.num_class,
args.imb_ratio_l)
U_SAMPLES_PER_CLASS = make_imb_data(args.ratio * args.num_max,
args.num_class, args.imb_ratio_u)
N_SAMPLES_PER_CLASS_T = torch.Tensor(N_SAMPLES_PER_CLASS)
print(args.out)
if args.dataset == 'cifar10':
print(f'==> Preparing imbalanced CIFAR-10')
train_labeled_set, train_unlabeled_set, test_set = get_cifar10(
'/home/jaehyung/data', N_SAMPLES_PER_CLASS, U_SAMPLES_PER_CLASS,
args.out)
elif args.dataset == 'stl10':
print(f'==> Preparing imbalanced STL-10')
train_labeled_set, train_unlabeled_set, test_set = get_stl10(
'/home/jaehyung/data', N_SAMPLES_PER_CLASS, args.out)
elif args.dataset == 'cifar100':
print(f'==> Preparing imbalanced CIFAR-100')
train_labeled_set, train_unlabeled_set, test_set = get_cifar100(
'/home/jaehyung/data', N_SAMPLES_PER_CLASS, U_SAMPLES_PER_CLASS,
args.out)
labeled_trainloader = data.DataLoader(train_labeled_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
drop_last=True)
unlabeled_trainloader = data.DataLoader(train_unlabeled_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
drop_last=True)
test_loader = data.DataLoader(test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
# Model
print("==> creating WRN-28-2")
def create_model(ema=False):
model = models.WRN(2, args.num_class)
model = model.cuda()
if ema:
for param in model.parameters():
param.detach_()
return model
model = create_model()
ema_model = create_model(ema=True)
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
train_criterion = SemiLoss()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
ema_optimizer = WeightEMA(model,
ema_model,
lr=args.lr,
alpha=args.ema_decay)
start_epoch = 0
# Resume
title = 'Imbalanced' + '-' + args.dataset + '-' + args.semi_method
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isfile(
args.resume), 'Error: no checkpoint directory found!'
args.out = os.path.dirname(args.resume)
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
ema_model.load_state_dict(checkpoint['ema_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(args.out, 'log.txt'),
title=title,
resume=True)
else:
logger = Logger(os.path.join(args.out, 'log.txt'), title=title)
logger.set_names([
'Train Loss', 'Train Loss X', 'Train Loss U', 'Test Loss',
'Test Acc.', 'Test GM.'
])
test_accs = []
test_gms = []
# Default values for MixMatch and DARP
emp_distb_u = torch.ones(args.num_class) / args.num_class
pseudo_orig = torch.ones(len(train_unlabeled_set.data),
args.num_class) / args.num_class
pseudo_refine = torch.ones(len(train_unlabeled_set.data),
args.num_class) / args.num_class
# Main function
for epoch in range(start_epoch, args.epochs):
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.epochs, state['lr']))
# Use the estimated distribution of unlabeled data
if args.est:
if args.dataset == 'cifar10':
est_name = './estimation/cifar10@N_1500_r_{}_{}_estim.npy'.format(
args.imb_ratio_l, args.imb_ratio_u)
else:
est_name = './estimation/stl10@N_450_r_{}_estim.npy'.format(
args.imb_ratio_l)
est_disb = np.load(est_name)
target_disb = len(train_unlabeled_set.data) * torch.Tensor(
est_disb) / np.sum(est_disb)
# Use the inferred distribution with labeled data
else:
target_disb = N_SAMPLES_PER_CLASS_T * len(
train_unlabeled_set.data) / sum(N_SAMPLES_PER_CLASS)
train_loss, train_loss_x, train_loss_u, emp_distb_u, pseudo_orig, pseudo_refine = trains(
args, labeled_trainloader, unlabeled_trainloader, model, optimizer,
ema_optimizer, train_criterion, epoch, use_cuda, target_disb,
emp_distb_u, pseudo_orig, pseudo_refine)
# Evaluation part
test_loss, test_acc, test_cls, test_gm = validate(
test_loader,
ema_model,
criterion,
use_cuda,
mode='Test Stats',
num_class=args.num_class)
# Append logger file
logger.append([
train_loss, train_loss_x, train_loss_u, test_loss, test_acc,
test_gm
])
# Save models
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'ema_state_dict': ema_model.state_dict(),
'optimizer': optimizer.state_dict(),
}, epoch + 1, args.out)
test_accs.append(test_acc)
test_gms.append(test_gm)
logger.close()
# Print the final results
print('Mean bAcc:')
print(np.mean(test_accs[-20:]))
print('Mean GM:')
print(np.mean(test_gms[-20:]))
print('Name of saved folder:')
print(args.out)
print(result[3])
# if step % 10 == 0:
# examples_per_sec = batch_size/duration
# sec_per_batch = float(duration)
# format_str = '%s: step %d, loss = %.2f (%.1f examples/sec; %.3f sec/batch)'
# print(format_str % (datetime.now(), step, result[1], examples_per_sec, sec_per_batch))
#
# if step % 1000 == 0:
# print("%s: step %d, evaluating test set" % (datetime.now(), step))
# correct_count = 0
# num_tst_examples = tst[0].shape[0]
# for tst_idx in range(0, num_tst_examples, batch_size):
# X_tst = tst[0][tst_idx:np.min([tst_idx+batch_size, num_tst_examples]), :]
# X_tst = X_tst.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
# Y_tst = tst[1][tst_idx:np.min([tst_idx+batch_size, num_tst_examples])]
# correct_count += total_correct.eval({
# raw_images: X_tst,
# labels: Y_tst,
# dropout_keep_prob: 1.0
# })
# print("%s tst accuracy is = %s" % (datetime.now(), float(correct_count)/num_tst_examples))
if __name__ == '__main__':
import dataset
batch_size = 20
trn_generator, val_generator = dataset.get_cifar10(batch_size=batch_size)
train(trn_generator, val_generator, steps_per_epoch=50000/batch_size, batch_size=batch_size)
# if step % 10 == 0:
# examples_per_sec = batch_size/duration
# sec_per_batch = float(duration)
# format_str = '%s: step %d, loss = %.2f (%.1f examples/sec; %.3f sec/batch)'
# print(format_str % (datetime.now(), step, result[1], examples_per_sec, sec_per_batch))
#
# if step % 1000 == 0:
# print("%s: step %d, evaluating test set" % (datetime.now(), step))
# correct_count = 0
# num_tst_examples = tst[0].shape[0]
# for tst_idx in range(0, num_tst_examples, batch_size):
# X_tst = tst[0][tst_idx:np.min([tst_idx+batch_size, num_tst_examples]), :]
# X_tst = X_tst.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
# Y_tst = tst[1][tst_idx:np.min([tst_idx+batch_size, num_tst_examples])]
# correct_count += total_correct.eval({
# raw_images: X_tst,
# labels: Y_tst,
# dropout_keep_prob: 1.0
# })
# print("%s tst accuracy is = %s" % (datetime.now(), float(correct_count)/num_tst_examples))
if __name__ == '__main__':
import dataset
batch_size = 20
trn_generator, val_generator = dataset.get_cifar10(batch_size=batch_size)
train(trn_generator,
val_generator,
steps_per_epoch=50000 / batch_size,
batch_size=batch_size)
from vgg import inference_vgg, batch_size
import tensorflow as tf
import numpy as np
def predict(img):
X = np.zeros((batch_size, 32, 32, 3))
X[0] = img
with tf.Graph().as_default():
in_images = tf.placeholder("float", [batch_size, 32, 32, 3])
images = tf.image.resize_images(in_images, 64, 64)
inference_op = inference_vgg(images,
dropout_keep_prob=tf.constant(
1.0, dtype=tf.float32),
input_shape=64)
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, "checkpoints/model.ckpt")
Y = sess.run(inference_op, feed_dict={in_images: X})
return Y[0]
if __name__ == '__main__':
import dataset
trn, tst = dataset.get_cifar10(10)
d = tst[11][1].reshape(3, 32, 32).transpose(1, 2, 0)
print predict(d)
from vgg import inference_vgg, batch_size
import tensorflow as tf
import numpy as np
def predict(img):
X = np.zeros((batch_size, 32,32, 3))
X[0] = img
with tf.Graph().as_default():
in_images = tf.placeholder("float", [batch_size, 32, 32, 3])
images = tf.image.resize_images(in_images, 64, 64)
inference_op = inference_vgg(images, dropout_keep_prob=tf.constant(1.0, dtype=tf.float32), input_shape=64)
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, "checkpoints/model.ckpt")
Y = sess.run(inference_op, feed_dict={in_images: X})
return Y[0]
if __name__ == '__main__':
import dataset
trn, tst = dataset.get_cifar10(10)
d = tst[11][1].reshape(3,32,32).transpose(1,2,0)
print predict(d)
def train(lr=0.0001, max_step=5000*10):
"""
train model
:param lr: This is the learning rate
"""
with tf.Graph().as_default():
in_images = tf.placeholder("float", [batch_size, 32, 32, 3])
images = tf.image.resize_images(in_images, 64, 64)
labels = tf.placeholder("int32", [batch_size])
dropout_keep_prob = tf.placeholder("float")
# Build a Graph that computes the logits predictions from the
# inference model.
# last_layer = inference_vgg(images, dropout_keep_prob)
last_layer = inference_vgg(images, dropout_keep_prob )
# Add a simple objective so we can calculate the backward pass.
objective = loss(last_layer, labels)
_, total_correct = evaluate(last_layer, labels)
optimizer = tf.train.RMSPropOptimizer(lr, 0.9)
global_step = tf.Variable(0, name="global_step", trainable=False)
train_step = optimizer.minimize(objective, global_step=global_step)
ema = tf.train.ExponentialMovingAverage(0.999)
maintain_averages_op = ema.apply([objective])
# grab summary variables we want to log
tf.scalar_summary("loss function", objective)
# tf.scalar_summary("accuracy", accuracy)
tf.scalar_summary("avg loss function", ema.average(objective))
# Create a saver.
saver = tf.train.Saver(tf.all_variables())
summary_op = tf.merge_all_summaries()
# Build an initialization operation.
initializer = tf.initialize_all_variables()
# Start running operations on the Graph.
with tf.Session() as sess:
sess.run(initializer)
writer = tf.train.SummaryWriter("train_logs", graph_def=sess.graph_def)
trn, tst = dataset.get_cifar10(batch_size)
for step in range(max_step):
# get batch and format data
batch = trn.next()
X = np.vstack(batch[0]).reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
Y = np.array(batch[1])
t0 = time.time()
result = sess.run(
[train_step, objective, summary_op, maintain_averages_op],
feed_dict = {
in_images: X,
labels: Y,
dropout_keep_prob: 0.5
}
)
duration = time.time() - t0
if np.isnan(result[1]):
print("gradient vanished/exploded")
return
if step % 10 == 0:
examples_per_sec = batch_size/duration
sec_per_batch = float(duration)
format_str = '%s: step %d, loss = %.4f (%.1f examples/sec; %.3f sec/batch)'
print(format_str % (datetime.now(), step, result[1], examples_per_sec, sec_per_batch))
if step % 100 == 0:
writer.add_summary(result[2], step)
if step % 1000 == 0:
print("%s: step %d, evaluating test set" % (datetime.now(), step))
correct_count = 0
num_tst_examples = tst[0].shape[0]
for tst_idx in range(0, num_tst_examples, batch_size):
X_tst = tst[0][tst_idx:np.min([tst_idx+batch_size, num_tst_examples]), :]
X_tst = X_tst.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
Y_tst = tst[1][tst_idx:np.min([tst_idx+batch_size, num_tst_examples])]
correct_count += total_correct.eval({
in_images: X_tst,
labels: Y_tst,
dropout_keep_prob: 1.0
})
accuracy = float(correct_count)/num_tst_examples
print("%s tst accuracy = %.3f" % (datetime.now(), accuracy))
if accuracy > 0.9:
checkpoint_path = saver.save(sess, "checkpoints/model.ckpt")
print("saving model %s" % checkpoint_path)
def train(lr=0.0001, max_step=5000 * 10):
"""
train model
:param lr: This is the learning rate
"""
with tf.Graph().as_default():
in_images = tf.placeholder("float", [batch_size, 32, 32, 3])
images = tf.image.resize_images(in_images, 64, 64)
labels = tf.placeholder("int32", [batch_size])
dropout_keep_prob = tf.placeholder("float")
# Build a Graph that computes the logits predictions from the
# inference model.
# last_layer = inference_vgg(images, dropout_keep_prob)
last_layer = inference_vgg(images, dropout_keep_prob)
# Add a simple objective so we can calculate the backward pass.
objective = loss(last_layer, labels)
_, total_correct = evaluate(last_layer, labels)
optimizer = tf.train.RMSPropOptimizer(lr, 0.9)
global_step = tf.Variable(0, name="global_step", trainable=False)
train_step = optimizer.minimize(objective, global_step=global_step)
ema = tf.train.ExponentialMovingAverage(0.999)
maintain_averages_op = ema.apply([objective])
# grab summary variables we want to log
tf.scalar_summary("loss function", objective)
# tf.scalar_summary("accuracy", accuracy)
tf.scalar_summary("avg loss function", ema.average(objective))
# Create a saver.
saver = tf.train.Saver(tf.all_variables())
summary_op = tf.merge_all_summaries()
# Build an initialization operation.
initializer = tf.initialize_all_variables()
# Start running operations on the Graph.
with tf.Session() as sess:
sess.run(initializer)
writer = tf.train.SummaryWriter("train_logs",
graph_def=sess.graph_def)
trn, tst = dataset.get_cifar10(batch_size)
for step in range(max_step):
# get batch and format data
batch = trn.next()
X = np.vstack(batch[0]).reshape(-1, 3, 32,
32).transpose(0, 2, 3, 1)
Y = np.array(batch[1])
t0 = time.time()
result = sess.run(
[train_step, objective, summary_op, maintain_averages_op],
feed_dict={
in_images: X,
labels: Y,
dropout_keep_prob: 0.5
})
duration = time.time() - t0
if np.isnan(result[1]):
print("gradient vanished/exploded")
return
if step % 10 == 0:
examples_per_sec = batch_size / duration
sec_per_batch = float(duration)
format_str = '%s: step %d, loss = %.4f (%.1f examples/sec; %.3f sec/batch)'
print(format_str % (datetime.now(), step, result[1],
examples_per_sec, sec_per_batch))
if step % 100 == 0:
writer.add_summary(result[2], step)
if step % 1000 == 0:
print("%s: step %d, evaluating test set" %
(datetime.now(), step))
correct_count = 0
num_tst_examples = tst[0].shape[0]
for tst_idx in range(0, num_tst_examples, batch_size):
X_tst = tst[0][tst_idx:np.min(
[tst_idx + batch_size, num_tst_examples]), :]
X_tst = X_tst.reshape(-1, 3, 32,
32).transpose(0, 2, 3, 1)
Y_tst = tst[1][tst_idx:np.min(
[tst_idx + batch_size, num_tst_examples])]
correct_count += total_correct.eval({
in_images:
X_tst,
labels:
Y_tst,
dropout_keep_prob:
1.0
})
accuracy = float(correct_count) / num_tst_examples
print("%s tst accuracy = %.3f" %
(datetime.now(), accuracy))
if accuracy > 0.9:
checkpoint_path = saver.save(sess,
"checkpoints/model.ckpt")
print("saving model %s" % checkpoint_path)