def train_source(config, base_network, classifier_gnn, dset_loaders):
# define loss functions
criterion_gedge = nn.BCELoss(reduction='mean')
ce_criterion = nn.CrossEntropyLoss()
# configure optimizer
optimizer_config = config['optimizer']
parameter_list = base_network.get_parameters() +\
[{'params': classifier_gnn.parameters(), 'lr_mult': 10, 'decay_mult': 2}]
optimizer = optimizer_config['type'](parameter_list, **(optimizer_config['optim_params']))
# configure learning rates
param_lr = []
for param_group in optimizer.param_groups:
param_lr.append(param_group['lr'])
schedule_param = optimizer_config['lr_param']
# start train loop
base_network.train()
classifier_gnn.train()
len_train_source = len(dset_loaders["source"])
for i in range(config['source_iters']):
optimizer = utils.inv_lr_scheduler(optimizer, i, **schedule_param)
optimizer.zero_grad()
# get input data
if i % len_train_source == 0:
iter_source = iter(dset_loaders["source"])
batch_source = iter_source.next()
inputs_source, labels_source = batch_source['img'].to(DEVICE), batch_source['target'].to(DEVICE)
# make forward pass for encoder and mlp head
features_source, logits_mlp = base_network(inputs_source)
mlp_loss = ce_criterion(logits_mlp, labels_source)
# make forward pass for gnn head
logits_gnn, edge_sim = classifier_gnn(features_source)
gnn_loss = ce_criterion(logits_gnn, labels_source)
# compute edge loss
edge_gt, edge_mask = classifier_gnn.label2edge(labels_source.unsqueeze(dim=0))
edge_loss = criterion_gedge(edge_sim.masked_select(edge_mask), edge_gt.masked_select(edge_mask))
# total loss and backpropagation
loss = mlp_loss + config['lambda_node'] * gnn_loss + config['lambda_edge'] * edge_loss
loss.backward()
optimizer.step()
# printout train loss
if i % 20 == 0 or i == config['source_iters'] - 1:
log_str = 'Iters:(%4d/%d)\tMLP loss:%.4f\tGNN loss:%.4f\tEdge loss:%.4f' % (i,
config['source_iters'], mlp_loss.item(), gnn_loss.item(), edge_loss.item())
utils.write_logs(config, log_str)
# evaluate network every test_interval
if i % config['test_interval'] == config['test_interval'] - 1:
evaluate(i, config, base_network, classifier_gnn, dset_loaders['target_test'])
return base_network, classifier_gnn
def supervised_testing():
test_x, test_y = get_data_set("test")
num_test = test_x.shape[0]
test_data = tf.data.Dataset.from_tensor_slices((test_x, test_y))
test_data = test_data.map(test_parse, num_parallel_calls=8)
test_data = test_data.batch(1)
test_iter = test_data.make_initializable_iterator()
x_test, y_test = test_iter.get_next()
X = tf.placeholder(tf.float32, [None, height, width, 3], name='Input')
Y = tf.placeholder(tf.int32, [None, num_classes], name='Label')
drop_rate = tf.placeholder(tf.float32)
logits = TESnet(X, "TESnet", drop_rate, reuse=False)
pred = tf.nn.softmax(logits)
saver = tf.train.Saver()
# Evaluate Model
correct_pred = tf.equal(tf.argmax(pred,1), tf.argmax(Y,1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
with tf.Session() as sess:
# Initialize variables
sess.run(tf.global_variables_initializer())
# Restore weights of model
saver.restore(sess, sv_model_dir)
log = "\n========== Supervised Testing Begin ==========\n"
write_logs(logs_sv, log, False)
test_start = time.time()
avg_acc = 0
sess.run(test_iter.initializer)
for i in range(num_test):
batch_start = time.time()
bx, by = sess.run([x_test, y_test])
acc = sess.run(accuracy, feed_dict={X:bx, Y:by, drop_rate:0.0})
avg_acc += acc
log = "Time {:2.5f}, Image {:05d}, Testing Accuracy = {:0.4f}".format(time.time()-batch_start, i+1, acc)
write_logs(logs_sv, log, False)
log = "\nTesting Accuracy = {:0.4f}\n".format(avg_acc/num_test)
write_logs(logs_sv, log, False)
log = "\nSupervised Testing Time: {:2.5f}".format(time.time()-test_start)
write_logs(logs_sv, log, False)
sess.close()
def main(args):
# fix random seed
random.seed(args.seed)
os.environ['PYTHONHASHSEED'] = str(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
assert args.source not in args.target, 'Source domain can not be one of the target domains'
# create train configurations
config = utils.build_config(args)
# prepare data
dsets, dset_loaders = utils.build_data(config)
# set base network
net_config = config['encoder']
base_network = net_config["name"](**net_config["params"])
base_network = base_network.to(DEVICE)
print(base_network)
# set GNN classifier
classifier_gnn = graph_net.ClassifierGNN(in_features=base_network.bottleneck.out_features,
edge_features=config['edge_features'],
nclasses=base_network.fc.out_features,
device=DEVICE)
classifier_gnn = classifier_gnn.to(DEVICE)
print(classifier_gnn)
# train on source domain and compute domain inheritability
log_str = '==> Step 1: Pre-training on the source dataset ...'
utils.write_logs(config, log_str)
base_network, classifier_gnn = trainer.train_source(config, base_network, classifier_gnn, dset_loaders)
log_str = '==> Finished pre-training on source!\n'
utils.write_logs(config, log_str)
log_str = '==> Step 2: Curriculum learning ...'
utils.write_logs(config, log_str)
######## Stage 1: find the closest target domain ##########
temp_test_loaders = dict(dset_loaders['target_test'])
max_inherit_domain = trainer.select_closest_domain(config, base_network, classifier_gnn, temp_test_loaders)
# iterate over all domains
for _ in range(len(config['data']['target']['name'])):
log_str = '==> Starting the adaptation on {} ...'.format(max_inherit_domain)
utils.write_logs(config, log_str)
######## Stage 2: adapt to the chosen target domain having the maximum inheritance/similarity ##########
base_network, classifier_gnn = trainer.adapt_target(config, base_network, classifier_gnn,
dset_loaders, max_inherit_domain)
log_str = '==> Finishing the adaptation on {}!\n'.format(max_inherit_domain)
utils.write_logs(config, log_str)
######### Stage 3: obtain the target pseudo labels and upgrade source domain ##########
trainer.upgrade_source_domain(config, max_inherit_domain, dsets,
dset_loaders, base_network, classifier_gnn)
######### Sage 1: recompute target domain inheritability/similarity ###########
# remove already considered domain
del temp_test_loaders[max_inherit_domain]
# find the maximum inheritability/similarity domain
if len(temp_test_loaders.keys()) > 0:
max_inherit_domain = trainer.select_closest_domain(config, base_network,
classifier_gnn, temp_test_loaders)
######### Step 3: fine-tuning stage ###########
log_str = '==> Step 3: Fine-tuning on pseudo-source dataset ...'
utils.write_logs(config, log_str)
config['source_iters'] = config['finetune_iters']
base_network, classifier_gnn = trainer.train_source(config, base_network, classifier_gnn, dset_loaders)
log_str = 'Finished training and evaluation!'
utils.write_logs(config, log_str)
# save models
if args.save_models:
torch.save(base_network.cpu().state_dict(), os.path.join(config['output_path'], 'base_network.pth.tar'))
torch.save(classifier_gnn.cpu().state_dict(), os.path.join(config['output_path'], 'classifier_gnn.pth.tar'))
def train(self):
noisyPatches, gtPatches = mutils.get_training_patches(
self.data_dir, self.dataset, self.patch_size, self.overlap_size)
# subtract the mean from bothdata
SL = np.reshape(noisyPatches,
[noisyPatches.shape[0], self.patch_size**2])
mSL = np.expand_dims(np.mean(SL, axis=1), axis=1)
rmSL = np.tile(mSL, [1, SL.shape[1]])
SL = SL - rmSL
SH = np.reshape(gtPatches, [gtPatches.shape[0], self.patch_size**2])
mSH = np.expand_dims(np.mean(SH, axis=1), axis=1)
rmSH = np.tile(mSH, [1, SH.shape[1]])
SH = SH - rmSH
# make sure patch is stored in column vector
SL = np.transpose(SL, [1, 0])
SH = np.transpose(SH, [1, 0])
hDim = SH.shape[0]
lDim = SL.shape[0]
# normalize signals
SL = cmod.normalise(SL)
SH = cmod.normalise(SH)
# Joint training
S = np.concatenate((SL, SH), axis=0)
# S = cmod.normalise(S)
# initialize dictionary
np.random.seed(1133221)
DL = np.random.randn(SL.shape[0], self.dict_size)
DH = np.random.randn(SH.shape[0], self.dict_size)
# join learning dict
D0 = np.concatenate((DL, DH), axis=0)
# X and D update options
lmbda = 0.1
optx = bpdn.BPDN.Options({
'Verbose': False,
'MaxMainIter': 1,
'rho': 50.0 * lmbda + 0.5
})
optd = cmod.CnstrMOD.Options({
'Verbose': False,
'MaxMainIter': 1,
'rho': S.shape[1] / 200.0
})
# update D update options
optd.update({'Y0': D0, 'U0': np.zeros((S.shape[0], D0.shape[1]))})
# create X update object
xstep = bpdn.BPDN(D0, S, lmbda, optx)
# create D update object
dstep = cmod.CnstrMOD(None, S, (D0.shape[1], S.shape[1]), optd)
# create dict learn object
opt = dictlrn.DictLearn.Options({'Verbose': True, 'MaxMainIter': 2})
d = dictlrn.DictLearn(xstep=xstep, dstep=dstep, opt=opt)
Dmx = d.solve()
print("DictLearn solv time: %.2fs" % d.timer.elapsed('solve'))
# get dictionary
DL1 = Dmx[0:self.patch_size**2, :]
DL1 = DL1.reshape(self.patch_size, self.patch_size, DL1.shape[1])
DH1 = Dmx[self.patch_size**2:, :]
DH1 = DH1.reshape(self.patch_size, self.patch_size, DH1.shape[1])
itsx = xstep.getitstat()
itsd = dstep.getitstat()
# write logs
mutils.write_logs(itsx, itsd, d.timer.elapsed('solve'),
self.output_dir, 'logs_%d_5' % self.dict_size)
# write dictionary
mutils.write_dicts(DL1, DH1, self.output_dir,
'dicts_%d_5' % self.dict_size)
parser.add_argument('--data_dir', type=str, default='')
parser.add_argument('--exp_dir', type=str, default='./test/')
parser.add_argument('--ckpt_file', type=str, default='')
parser.add_argument('--device', type=str, default='cuda')
# parser.add_argument('--multi_gpu', action='store_true')
parser.add_argument('--dataset', type=str, default='Simp')
parser.add_argument('--test_split', type=float, default=0.2)
parser.add_argument('--red_rate', type=float, default=0.0)
parser.add_argument('--validation_split', type=float, default=0.0)
parser.add_argument('--d_latent', type=int, default=256)
parser.add_argument('--batch_size', type=int, default=32)
parser.add_argument('--n_epochs', type=int, default=2000)
parser.add_argument('--size', type=int, default=256)
parser.add_argument('--logfile', type=str, default='test.txt')
args = parser.parse_args()
if args.device == 'cuda' and torch.cuda.is_available():
from subprocess import call
print('available gpus:')
call([
"nvidia-smi", "--format=csv",
"--query-gpu=index,name,driver_version,memory.total,memory.used,memory.free"
])
cudnn.benchmark = True
else:
args.device = 'cpu'
utils.prepare_directory(args.exp_dir)
utils.write_logs(args)
configure(args.exp_dir)
main(args)
def adapt_target(config, base_network, classifier_gnn, dset_loaders,
max_inherit_domain):
# define loss functions
criterion_gedge = nn.BCELoss(reduction='mean')
ce_criterion = nn.CrossEntropyLoss()
# add random layer and adversarial network
class_num = config['encoder']['params']['class_num']
random_layer = networks.RandomLayer([base_network.output_num(), class_num],
config['random_dim'], DEVICE)
adv_net = networks.AdversarialNetwork(config['random_dim'],
config['random_dim'],
config['ndomains'])
random_layer.to(DEVICE)
adv_net = adv_net.to(DEVICE)
# configure optimizer
optimizer_config = config['optimizer']
parameter_list = base_network.get_parameters() + adv_net.get_parameters() \
+ [{'params': classifier_gnn.parameters(), 'lr_mult': 10, 'decay_mult': 2}]
optimizer = optimizer_config['type'](parameter_list,
**(optimizer_config['optim_params']))
# configure learning rates
param_lr = []
for param_group in optimizer.param_groups:
param_lr.append(param_group['lr'])
schedule_param = optimizer_config['lr_param']
# start train loop
len_train_source = len(dset_loaders['source'])
len_train_target = len(dset_loaders['target_train'][max_inherit_domain])
# set nets in train mode
base_network.train()
classifier_gnn.train()
adv_net.train()
random_layer.train()
for i in range(config['adapt_iters']):
optimizer = utils.inv_lr_scheduler(optimizer, i, **schedule_param)
optimizer.zero_grad()
# get input data
if i % len_train_source == 0:
iter_source = iter(dset_loaders['source'])
if i % len_train_target == 0:
iter_target = iter(
dset_loaders['target_train'][max_inherit_domain])
batch_source = iter_source.next()
batch_target = iter_target.next()
inputs_source, inputs_target = batch_source['img'].to(
DEVICE), batch_target['img'].to(DEVICE)
labels_source = batch_source['target'].to(DEVICE)
domain_source, domain_target = batch_source['domain'].to(
DEVICE), batch_target['domain'].to(DEVICE)
domain_input = torch.cat([domain_source, domain_target], dim=0)
# make forward pass for encoder and mlp head
features_source, logits_mlp_source = base_network(inputs_source)
features_target, logits_mlp_target = base_network(inputs_target)
features = torch.cat((features_source, features_target), dim=0)
logits_mlp = torch.cat((logits_mlp_source, logits_mlp_target), dim=0)
softmax_mlp = nn.Softmax(dim=1)(logits_mlp)
mlp_loss = ce_criterion(logits_mlp_source, labels_source)
# *** GNN at work ***
# make forward pass for gnn head
logits_gnn, edge_sim = classifier_gnn(features)
gnn_loss = ce_criterion(logits_gnn[:labels_source.size(0)],
labels_source)
# compute pseudo-labels for affinity matrix by mlp classifier
out_target_class = torch.softmax(logits_mlp_target, dim=1)
target_score, target_pseudo_labels = out_target_class.max(1,
keepdim=True)
idx_pseudo = target_score > config['threshold']
target_pseudo_labels[~idx_pseudo] = classifier_gnn.mask_val
# combine source labels and target pseudo labels for edge_net
node_labels = torch.cat(
(labels_source, target_pseudo_labels.squeeze(dim=1)),
dim=0).unsqueeze(dim=0)
# compute source-target mask and ground truth for edge_net
edge_gt, edge_mask = classifier_gnn.label2edge(node_labels)
# compute edge loss
edge_loss = criterion_gedge(edge_sim.masked_select(edge_mask),
edge_gt.masked_select(edge_mask))
# *** Adversarial net at work ***
if config['method'] == 'CDAN+E':
entropy = transfer_loss.Entropy(softmax_mlp)
trans_loss = transfer_loss.CDAN(config['ndomains'],
[features, softmax_mlp], adv_net,
entropy, networks.calc_coeff(i),
random_layer, domain_input)
elif config['method'] == 'CDAN':
trans_loss = transfer_loss.CDAN(config['ndomains'],
[features, softmax_mlp], adv_net,
None, None, random_layer,
domain_input)
else:
raise ValueError('Method cannot be recognized.')
# total loss and backpropagation
loss = config['lambda_adv'] * trans_loss + mlp_loss +\
config['lambda_node'] * gnn_loss + config['lambda_edge'] * edge_loss
loss.backward()
optimizer.step()
# printout train loss
if i % 20 == 0 or i == config['adapt_iters'] - 1:
log_str = 'Iters:(%4d/%d)\tMLP loss: %.4f\t GNN Loss: %.4f\t Edge Loss: %.4f\t Transfer loss:%.4f' % (
i, config["adapt_iters"], mlp_loss.item(),
config['lambda_node'] * gnn_loss.item(), config['lambda_edge']
* edge_loss.item(), config['lambda_adv'] * trans_loss.item())
utils.write_logs(config, log_str)
# evaluate network every test_interval
if i % config['test_interval'] == config['test_interval'] - 1:
evaluate(i, config, base_network, classifier_gnn,
dset_loaders['target_test'])
return base_network, classifier_gnn
def main(args):
# fix random seed
random.seed(args.seed)
os.environ['PYTHONHASHSEED'] = str(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
# create train configurations
args.use_cgct_mask = True # used in CGCT for pseudo label mask in target datasets
config = utils.build_config(args)
# prepare data
dsets, dset_loaders = utils.build_data(config)
# set base network
net_config = config['encoder']
base_network = net_config["name"](**net_config["params"])
base_network = base_network.to(DEVICE)
print(base_network)
# set GNN classifier
classifier_gnn = graph_net.ClassifierGNN(
in_features=base_network.bottleneck.out_features,
edge_features=config['edge_features'],
nclasses=base_network.fc.out_features,
device=DEVICE)
classifier_gnn = classifier_gnn.to(DEVICE)
print(classifier_gnn)
# train on source domain
log_str = '==> Step 1: Pre-training on the source dataset ...'
utils.write_logs(config, log_str)
base_network, classifier_gnn = trainer.train_source(
config, base_network, classifier_gnn, dset_loaders)
log_str = '==> Finished pre-training on source!\n'
utils.write_logs(config, log_str)
# create random layer and adversarial network
class_num = config['encoder']['params']['class_num']
random_layer = networks.RandomLayer([base_network.output_num(), class_num],
config['random_dim'], DEVICE)
adv_net = networks.AdversarialNetwork(config['random_dim'],
config['random_dim'],
config['ndomains'])
random_layer = random_layer.to(DEVICE)
adv_net = adv_net.to(DEVICE)
print(random_layer)
print(adv_net)
# run adaptation episodes
log_str = '==> Starting the adaptation'
utils.write_logs(config, log_str)
for curri_iter in range(len(config['data']['target']['name'])):
######## Step 1: train one adaptation episod on combined target domains ##########
target_train_datasets = preprocess.ConcatDataset(
dsets['target_train'].values())
dset_loaders['target_train'] = DataLoader(
dataset=target_train_datasets,
batch_size=config['data']['target']['batch_size'],
shuffle=True,
num_workers=config['num_workers'],
drop_last=True)
base_network, classifier_gnn = trainer.adapt_target_cgct(
config, base_network, classifier_gnn, dset_loaders, random_layer,
adv_net)
log_str = '==> Finishing {} adaptation episode!\n'.format(curri_iter)
utils.write_logs(config, log_str)
######### Step 2: obtain the target pseudo labels and upgrade target domains ##########
trainer.upgrade_target_domains(config, dsets, dset_loaders,
base_network, classifier_gnn,
curri_iter)
######### Step 3: fine-tuning stage ###########
log_str = '==> Step 3: Fine-tuning on pseudo-source dataset ...'
utils.write_logs(config, log_str)
config['source_iters'] = config['finetune_iters']
base_network, classifier_gnn = trainer.train_source(
config, base_network, classifier_gnn, dset_loaders)
log_str = 'Finished training and evaluation!'
utils.write_logs(config, log_str)
# save models
if args.save_models:
torch.save(base_network.cpu().state_dict(),
os.path.join(config['output_path'], 'base_network.pth.tar'))
torch.save(
classifier_gnn.cpu().state_dict(),
os.path.join(config['output_path'], 'classifier_gnn.pth.tar'))
def supervised_training():
train_x, train_y = get_data_set("train")
with open('..\\data\\svtrain.p', 'rb') as fp:
idx = pickle.load(fp)
train_x = train_x[idx, :]
train_y = train_y[idx, :]
num_train = train_x.shape[0]
train_data = tf.data.Dataset.from_tensor_slices((train_x, train_y))
train_data = train_data.shuffle(num_train)
train_data = train_data.map(train_parse, num_parallel_calls=8)
train_data = train_data.batch(batch_size)
train_iter = train_data.make_initializable_iterator()
x_train, y_train = train_iter.get_next()
test_x, test_y = get_data_set("test")
num_test = test_x.shape[0]
test_data = tf.data.Dataset.from_tensor_slices((test_x, test_y))
test_data = test_data.map(test_parse, num_parallel_calls=8)
test_data = test_data.batch(1)
test_iter = test_data.make_initializable_iterator()
x_test, y_test = test_iter.get_next()
X = tf.placeholder(tf.float32, [None, height, width, 3], name='Input')
Y = tf.placeholder(tf.int32, [None, num_classes], name='Label')
drop_rate = tf.placeholder(tf.float32)
logits = TESnet(X, "TESnet", drop_rate, reuse=False)
pred = tf.nn.softmax(logits)
# Learning Rate
with tf.variable_scope('learning_rate'):
lr_v = tf.Variable(lr_init, trainable=False)
# Loss Function
with tf.name_scope("Cross_Entropy_Loss"):
loss_op = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=Y))
sum_loss_op = tf.summary.scalar("Cross_Entropy_Loss", loss_op)
# Optimizer
optimizer = tf.train.AdamOptimizer(lr_v)
gvs = optimizer.compute_gradients(loss_op)
capped_gvs = [(tf.clip_by_value(grad,-1.0, 1.0), var) for grad, var in gvs]
train_op = optimizer.apply_gradients(capped_gvs)
saver = tf.train.Saver()
# Evaluate Model
correct_pred = tf.equal(tf.argmax(pred,1), tf.argmax(Y,1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
sum_acc_op = tf.summary.scalar("Accuracy", accuracy)
num_batches = int(math.ceil(num_train/batch_size))
with tf.Session() as sess:
log = "\n========== Supervised Training Begin ==========\n"
write_logs(logs_sv, log, True)
train_start = time.time()
# Initialize variables
sess.run(tf.global_variables_initializer())
# Op to write logs to Tensorboard
train_sum_writer = tf.summary.FileWriter(logs_dir, tf.get_default_graph())
for epoch in range(num_epoches):
epoch_start = time.time()
if (epoch == 70):
new_lr = lr_v * lr_decay
sess.run(tf.assign(lr_v, new_lr))
log = "** New learning rate: %1.9f **\n" % (lr_v.eval())
write_logs(logs_sv, log, False)
elif epoch == 0:
sess.run(tf.assign(lr_v, lr_init))
log = "** Initial learning rate: %1.9f **\n" % (lr_init)
write_logs(logs_sv, log, False)
avg_loss = 0
avg_acc = 0
sess.run(train_iter.initializer)
for batch in range(num_batches):
batch_start = time.time()
bx, by = sess.run([x_train, y_train])
sess.run([train_op], feed_dict={X:bx, Y:by, drop_rate:dropout})
loss, acc, sum_loss, sum_acc = sess.run([loss_op, accuracy, sum_loss_op, sum_acc_op], feed_dict={X:bx, Y:by, drop_rate:0.0})
avg_loss += loss
avg_acc += acc
train_sum_writer.add_summary(sum_loss, epoch*num_batches+batch)
train_sum_writer.add_summary(sum_acc, epoch*num_batches+batch)
log = "Time {:2.5f}, Epoch {}, Batch {}, Loss = {:2.5f}, Training Accuracy = {:0.4f}".format(time.time()-batch_start, epoch, batch, loss, acc)
write_logs(logs_sv, log, False)
log = "\nTime {:2.5f}, Epoch {}, Average Loss = {:2.5f}, Training Average Accuracy = {:0.4f}\n"\
.format(time.time()-epoch_start, epoch, avg_loss/num_batches, avg_acc/num_batches)
write_logs(logs_sv, log, False)
log = "\nSupervised Training Time: {:2.5f}".format(time.time()-train_start)
write_logs(logs_sv, log, False)
log = "\n========== Supervised Training End ==========\n"
write_logs(logs_sv, log, False)
# Save model
save_path = saver.save(sess, sv_model_dir)
log = "Model is saved in file: %s" % save_path
write_logs(logs_sv, log, False)
log = "\n========== Supervised Testing Begin ==========\n"
write_logs(logs_sv, log, False)
test_start = time.time()
avg_acc = 0
sess.run(test_iter.initializer)
for i in range(num_test):
batch_start = time.time()
bx, by = sess.run([x_test, y_test])
acc = sess.run(accuracy, feed_dict={X:bx, Y:by, drop_rate:0.0})
avg_acc += acc
log = "Time {:2.5f}, Image {:05d}, Testing Accuracy = {:0.4f}".format(time.time()-batch_start, i+1, acc)
write_logs(logs_sv, log, False)
log = "\nTesting Accuracy = {:0.4f}\n".format(avg_acc/num_test)
write_logs(logs_sv, log, False)
log = "\nSupervised Testing Time: {:2.5f}".format(time.time()-test_start)
write_logs(logs_sv, log, False)
sess.close()
type=str,
default='resnet-10',
choices=['resnet-10', 'resnet-18', 'resnet-34', 'resnet-50'])
parser.add_argument('--device',
type=str,
default='cuda',
choices=['cpu', 'cuda'])
parser.add_argument('--multi_gpu', type=bool, default=True)
parser.add_argument('--init_lr', type=float, default=0.1)
parser.add_argument('--lr_dec_rate', type=float, default=0.1)
parser.add_argument('--lr_dec_int', type=int, default=100)
parser.add_argument('--weight_decay', type=float, default=5e-4)
parser.add_argument('--momentum', type=float, default=0.9)
parser.add_argument('--tensorboard', type=bool, default=True)
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--test_batch_size', type=int, default=128)
parser.add_argument('--n_epochs', type=int, default=300)
FLAGS = parser.parse_args()
if FLAGS.device == 'cuda' and torch.cuda.is_available():
cudnn.benchmark = True
else:
FLAGS.device = 'cpu'
utils.write_logs(FLAGS)
if FLAGS.tensorboard:
from tensorboard_logger import configure
configure(FLAGS.exp_dir)
main(FLAGS)
def validate(ldr_dir, hdr_dir, gen_dir, logs_dir, img_height, img_width):
X = tf.placeholder(tf.float32, [1, img_height, img_width, 3])
Y = tf.placeholder(tf.float32, [1, img_height, img_width, 3])
valid_ldr_path, _ = get_filepath(ldr_dir, '.png')
valid_hdr_path, valid_hdr_name = get_filepath(hdr_dir, '.hdr')
num_valid = len(valid_hdr_path)
# Data loader
dataset = tf.data.Dataset.from_tensor_slices(
(valid_ldr_path, valid_hdr_path))
dataset = dataset.map(valid_parse, num_parallel_calls=4)
dataset = dataset.batch(1)
iter = dataset.make_one_shot_iterator()
ldr_img, hdr_img, Hth = iter.get_next()
alpha = alpha_msk(X)
# Prediction
with tf.name_scope('HDR_CNN'):
hdr_nn = hdrcnn(X, is_training=False, reuse=False)
hdr_final = get_final_hdr(X, hdr_nn)
# Loss functions
with tf.name_scope('Loss'):
irloss, dirloss = loss_fn(X, hdr_nn, Y)
saver = tf.train.Saver(tf.global_variables())
with tf.Session() as sess:
# Initialize variables
sess.run(tf.global_variables_initializer())
# Restore weights of model
saver.restore(sess, model_dir)
# Validation
log = "\n========== Validation Begin ==========\n"
write_logs(logs_dir, log, True)
valid_start = time.time()
avg_irloss = 0
avg_dirloss = 0
for f in valid_hdr_name:
valid_img_start = time.time()
ldr_image, hdr_image, Hth_val = sess.run([ldr_img, hdr_img, Hth])
alpha_val, hdr_pred, irloss_val, dirloss_val = sess.run(
[alpha, hdr_final, irloss, dirloss],
feed_dict={
X: ldr_image,
Y: hdr_image
})
avg_irloss += irloss_val
avg_dirloss += dirloss_val
f1, _ = f.split("_")
img_write(gen_dir, 'alpha_' + f1 + '_HDR.png', alpha_val, 'PNG-FI')
# Gamma correction
hdr_pred_save = np.multiply(Hth_val, np.maximum(hdr_pred, 0.0))
img_write(gen_dir, 'pred_' + f, hdr_pred_save, 'HDR-FI')
# Tone mapping
hdr_pred_gamma = np.power(np.maximum(hdr_pred, 0.0), gamma)
ldr_tone = reinhard02(hdr_pred_gamma, a=0.18)
img_write(gen_dir, 'tm_' + f1 + '_HDR.png', ldr_tone, 'PNG-FI')
log = "Image {}, Time {:2.5f}, Shape = {}, I/R Loss = {:2.5f}, Direct Loss = {:2.5f}".format(
f,
time.time() - valid_img_start, hdr_pred.shape, irloss_val,
dirloss_val)
write_logs(logs_dir, log, False)
log = "\nAverage I/R Loss = {:2.5f}, Average Direct Loss = {:2.5f}".format(
avg_irloss / num_valid, avg_dirloss / num_valid)
write_logs(logs_dir, log, False)
log = "\nValidation Time: {:2.5f}".format(time.time() - valid_start)
write_logs(logs_dir, log, False)
log = "\n========== Validation End ==========\n"
write_logs(logs_dir, log, False)
sess.close()
def train_and_evaluate():
svhn, mnist = get_datasets(is_training=True)
source_dataset = svhn if SOURCE_DATA == 'svhn' else mnist
target_dataset = mnist if SOURCE_DATA == 'svhn' else svhn
weights = make_weights_for_balanced_classes(source_dataset, num_classes=10)
sampler = WeightedRandomSampler(weights, len(weights))
source_loader = DataLoader(source_dataset,
BATCH_SIZE,
sampler=sampler,
pin_memory=True,
drop_last=True)
target_loader = DataLoader(target_dataset,
BATCH_SIZE,
shuffle=True,
pin_memory=True,
drop_last=True)
val_svhn, val_mnist = get_datasets(is_training=False)
val_svhn_loader = DataLoader(val_svhn,
BATCH_SIZE,
shuffle=False,
drop_last=False)
val_mnist_loader = DataLoader(val_mnist,
BATCH_SIZE,
shuffle=False,
drop_last=False)
print('\nsource dataset is', SOURCE_DATA, '\n')
num_steps_per_epoch = math.floor(min(len(svhn), len(mnist)) / BATCH_SIZE)
embedder = Network(image_size=(32, 32),
embedding_dim=EMBEDDING_DIM).to(DEVICE)
classifier = nn.Linear(EMBEDDING_DIM, 10).to(DEVICE)
model = nn.Sequential(embedder, classifier)
model.train()
optimizer = optim.Adam(lr=1e-3,
params=model.parameters(),
weight_decay=1e-3)
scheduler = CosineAnnealingLR(optimizer,
T_max=num_steps_per_epoch * NUM_EPOCHS -
DELAY,
eta_min=1e-6)
cross_entropy = nn.CrossEntropyLoss()
association = WalkerVisitLosses()
text = 'e:{0:2d}, i:{1:3d}, classification loss: {2:.3f}, ' +\
'walker loss: {3:.3f}, visit loss: {4:.4f}, ' +\
'total loss: {5:.3f}, lr: {6:.6f}'
logs, val_logs = [], []
i = 0 # iteration
for e in range(NUM_EPOCHS):
model.train()
for (x_source, y_source), (x_target, _) in zip(source_loader,
target_loader):
x_source = x_source.to(DEVICE)
x_target = x_target.to(DEVICE)
y_source = y_source.to(DEVICE)
x = torch.cat([x_source, x_target], dim=0)
embeddings = embedder(x)
a, b = torch.split(embeddings, BATCH_SIZE, dim=0)
logits = classifier(a)
usual_loss = cross_entropy(logits, y_source)
walker_loss, visit_loss = association(a, b, y_source)
if i > DELAY:
growth = torch.clamp(
torch.tensor((i - DELAY) / GROWTH_STEPS).to(DEVICE), 0.0,
1.0)
loss = usual_loss + growth * (BETA1 * walker_loss +
BETA2 * visit_loss)
else:
loss = usual_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i > DELAY:
scheduler.step()
lr = scheduler.get_lr()[0]
log = (e, i, usual_loss.item(), walker_loss.item(),
visit_loss.item(), loss.item(), lr)
print(text.format(*log))
logs.append(log)
i += 1
result1 = evaluate(model, cross_entropy, val_svhn_loader, DEVICE)
result2 = evaluate(model, cross_entropy, val_mnist_loader, DEVICE)
print('\nsvhn loss {0:.3f} and accuracy {1:.3f}'.format(*result1))
print('mnist loss {0:.3f} and accuracy {1:.3f}\n'.format(*result2))
val_logs.append((i, ) + result1 + result2)
torch.save(model.state_dict(), SAVE_PATH)
write_logs(logs, val_logs, LOGS_PATH)
def unsupervised_testing():
test_x, test_y = get_data_set("test")
num_test = test_x.shape[0]
test_data = tf.data.Dataset.from_tensor_slices((test_x, test_y))
test_data = test_data.map(test_parse, num_parallel_calls=8)
test_data = test_data.batch(1)
test_iter = test_data.make_initializable_iterator()
x_test, y_test = test_iter.get_next()
X = tf.placeholder(tf.float32, [None, height, width, 3], name='Input')
Y = tf.placeholder(tf.int32, [None, num_classes], name='Label')
drop_rate = tf.placeholder(tf.float32)
is_labeled = tf.placeholder(tf.bool)
# Networks
logits_std = TESnet(X, "Student", drop_rate, reuse=False, getter=None)
pred_std = tf.nn.softmax(logits_std)
logits_tc = TESnet(X, "Teacher", drop_rate, reuse=False, getter=None)
pred_tc = tf.nn.softmax(logits_tc)
# Evaluate Model
crt_pred_std = tf.equal(tf.argmax(pred_std, 1), tf.argmax(Y, 1))
acc_std = tf.reduce_mean(tf.cast(crt_pred_std, tf.float32))
sum_acc_std_op = tf.summary.scalar("Student Accuracy", acc_std)
crt_pred_tc = tf.equal(tf.argmax(pred_tc, 1), tf.argmax(Y, 1))
acc_tc = tf.reduce_mean(tf.cast(crt_pred_tc, tf.float32))
sum_acc_tc_op = tf.summary.scalar("Teacher Accuracy", acc_tc)
sum_acc_op = tf.summary.merge([sum_acc_std_op, sum_acc_tc_op])
saver = tf.train.Saver()
with tf.Session() as sess:
# Initialize variables
sess.run(tf.global_variables_initializer())
# Restore weights of model
saver.restore(sess, usv_model_dir)
log = "\n========== Semi-supervised Testing Begin ==========\n"
write_logs(logs_usv, log, False)
test_start = time.time()
avg_acc_std = 0
avg_acc_tc = 0
sess.run(test_iter.initializer)
for i in range(num_test):
batch_start = time.time()
bx, by = sess.run([x_test, y_test])
acc_std_val, acc_tc_val = sess.run([acc_std, acc_tc],
feed_dict={
X: bx,
Y: by,
drop_rate: 0.0
})
avg_acc_std += acc_std_val
avg_acc_tc += acc_tc_val
log = "Time {:2.5f}, Image {:05d}, Student Accuracy = {:0.4f}, Teacher Accuracy = {:0.4f}"\
.format(time.time()-batch_start, i+1, acc_std_val, acc_tc_val)
write_logs(logs_usv, log, False)
log = "\nTesting Student Accuracy = {:0.4f}, Testing Teacher Accuracy = {:0.4f}\n".format(
avg_acc_std / num_test, avg_acc_tc / num_test)
write_logs(logs_usv, log, False)
log = "\nSemi-supervised Testing Time: {:2.5f}".format(time.time() -
test_start)
write_logs(logs_usv, log, False)
sess.close()
def testing(ldr_dir, gen_dir, logs_dir, img_height, img_width, Hth):
X = tf.placeholder(tf.float32, [1, img_height , img_width, 3])
Y = tf.placeholder(tf.float32, [1, img_height, img_width, 3])
test_ldr_path, test_ldr_name = get_filepath(ldr_dir, '.png')
num_test = len(test_ldr_path)
test_ldr_path = tf.constant(test_ldr_path)
# Data loader
dataset = tf.data.Dataset.from_tensor_slices(test_ldr_path)
dataset = dataset.map(test_parse, num_parallel_calls=4)
#dataset = dataset.batch(1)
iter = dataset.make_one_shot_iterator()
ldr_img = iter.get_next()
alpha = alpha_msk(X)
# Prediction
with tf.name_scope('HDR_CNN'):
hdr_nn = hdrcnn(X, is_training=False, reuse=False)
hdr_final = get_final_hdr(X, hdr_nn)
saver = tf.train.Saver(tf.global_variables())
with tf.Session() as sess:
# Initialize variables
sess.run(tf.global_variables_initializer())
# Restore weights of model
saver.restore(sess, model_dir)
# Validation
log = "\n========== Test Begin ==========\n"
write_logs(logs_dir, log, True)
test_start = time.time()
avg_irloss = 0
avg_dirloss = 0
for f in test_ldr_name:
test_img_start = time.time()
ldr_image = sess.run([ldr_img])
alpha_val, hdr_pred = sess.run([alpha, hdr_final], feed_dict={X:ldr_image})
f1, _ = f.split("_")
img_write(gen_dir, 'alpha_'+f1+'_HDR.png', alpha_val, 'PNG-FI')
# Gamma correction
hdr_pred_save = np.multiply(Hth, np.maximum(hdr_pred, 0.0))
img_write(gen_dir, 'pred_'+f1+'_HDR.hdr', hdr_pred_save, 'HDR-FI')
# Tone mapping
hdr_pred_gamma = np.power(np.maximum(hdr_pred, 0.0), gamma)
ldr_tone = reinhard02(hdr_pred_gamma, a=0.18)
img_write(gen_dir, 'tm_'+f1+'_HDR.png', ldr_tone, 'PNG-FI')
log = "Image {}, Time {:2.5f}, Shape = {}".format(f, time.time()-test_img_start, hdr_pred.shape)
write_logs(logs_dir, log, False)
log = "\nTest Time: {:2.5f}".format(time.time()-test_start)
write_logs(logs_dir, log, False)
log = "\n========== Test End ==========\n"
write_logs(logs_dir, log, False)
sess.close()
