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()