def plot(ds, d_name, inds=[[1, 2, 3], [1, 4, 5]], add=0): d = os.path.join('output', d_name) check_dir_exists([d]) for i in range(len(inds)): imgs = [] for j in range(len(inds[i])): tmp = 0 if j == 0 else add imgs.append(ds[inds[i][j] + tmp][0]) save_image(imgs, os.path.join(d, 'img_%d.jpg' % (i)))
def generate_feature(output_file): ''' Generate feature of given dataset. 1. laod model 2. for each data in train_loader => (encoded, (label, img_name)) 3. output a dictionary {'features': [encoded], 'labels': [(label, img_name)]} :param output_file: output file name ''' model_name = 'model/%s_%s%s_model-%s.pkl' % (args.main_model, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) print(model_name) assert os.path.exists(model_name) print('Loading model ...') if cuda: mol = torch.load(model_name).to(device) else: mol = torch.load(model_name, map_location='cpu').to(device) # print(torchsummary.summary(autoencoder, input_size=(3, HEIGHT, WEIGHT))) train_loader = getDataLoader(args, kwargs) check_dir_exists(['feature']) features = [] labels = [] for step, (x, y) in enumerate(train_loader): b_x = Variable(x).cuda() if cuda else Variable(x) label = [(y[0][i], y[1][i]) for i in range(len(y[0]))] labels.extend(label) if args.main_model == 'AE': feature, _ = mol(b_x) elif args.main_model == 'VAE': _, mu, std = mol(b_x) feature = torch.cat([mu, std]) elif args.main_model == 'AEClass': feature, _, _ = mol(b_x) elif args.main_model == 'vgg_classifier': feature = mol.get_encode_features(b_x) f = feature.cpu() if cuda else feature f = f.data.view(b_x.shape[0], -1).numpy().tolist() features.extend(f) if step % 10 == 0: print('Step %d finished.' % step) out = {'features': features, 'labels': labels} with open(output_file, 'w') as f: json.dump(out, f) return out
def get_features(mol_short='inception_v3'): ################################################################ # Arguments ################################################################ ae_args = train_args() cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} # global ae_args, cuda, device, kwargs start_time = time.time() args = ae_args model_name = 'model/%s_%s%s_model-%s.pkl' % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % ( mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') mol = torch.load(model_name).to(device) else: print('Init model ...') mol = inception_v3_features(pretrained=True, training=False).to(device) # train_loader = getDataLoader(args, kwargs) test_loader = getDataLoader(args, kwargs, train='test') check_dir_exists(['res/', 'model', 'res/evaluation_pic', evaluation_dir]) t_per_img = [] for epoch in range(1): step_time = time.time() for step, (x, y) in enumerate(test_loader): b_x = Variable(x, volatile=True).cuda() if cuda else Variable(x) t0 = time.time() features = mol(b_x) t_tmp = (time.time() - t0) / len(b_x) * 1000 t_per_img.append(t_tmp) print( 'cost %.6fms per image this batch. cost %.6fms per image till now.' % (t_tmp, np.mean(sorted(t_per_img)[1:-1]))) print('Finished. Totally cost %.2f' % (time.time() - start_time))
def __init__(self, name='inception_finetune_experiments'): self.args = train_args() self._get_gpu() self.cuda = torch.cuda.is_available() and self.args.gpu != -2 self.device = torch.device("cuda" if self.cuda else "cpu") self.kwargs = {'num_workers': 6, 'pin_memory': True} self.mol_dir = 'model' self.log_dir = os.path.join('log') self.log_file = os.path.join(self.log_dir, 'log_%s.log' % name) self.summary_dir = os.path.join('summary', self.args.name) check_dir_exists([ self.mol_dir, self.log_dir, os.path.dirname(self.summary_dir), self.summary_dir ]) self.writer = SummaryWriter(self.summary_dir) logging.basicConfig(filename=self.log_file, level=logging.INFO)
def get_features(mol_short='resnet50'): ################################################################ # Arguments ################################################################ ae_args = train_args() cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} # global ae_args, cuda, device, kwargs start_time = time.time() args = ae_args model_name = 'model/%s_%s%s_model-%s.pkl' % ( mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % ( mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') mol = torch.load(model_name).to(device) else: print('Init model ...') mol = ResNet().to(device) test_loader = getDataLoader(args, kwargs, train='test') loss_class = nn.CrossEntropyLoss().cuda(cuda) check_dir_exists(['res/', 'model', 'res/evaluation_pic', evaluation_dir]) total, correct, top5correct, loss_total = 0, 0, 0, 0 t_per_img = [] for epoch in range(1): step_time = time.time() for step, (x, y) in enumerate(test_loader): t0 = time.time() b_x = Variable(x, volatile=True).cuda() if cuda else Variable(x) prob_class = mol.get_fc_features(b_x, return_both=False) t_tmp = (time.time() - t0) / len(b_x) * 1000 t_per_img.append(t_tmp) print('cost %.6fms per image this batch. cost %.6fms per image till now.' % ( t_tmp, np.mean(sorted(t_per_img)[1:-1])))
def train_cls(args): ################################################################ # Arguments ################################################################ cuda = args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} ################################################################ # Model name & Directory ################################################################ mol_short = get_mol_short(args.model) + '(cls)' model_name = 'model/%s_%s_model.pkl' % ( mol_short, args.model, ) print('[Model] model name is', model_name) pic_dir = 'res/%s_%s%s-%s/' % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % ( mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) log_dir = 'log/log_%s_%s%s_model-%s/' % \ (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) remove_dir_exists([log_dir]) check_dir_exists([ 'res/', 'model', pic_dir, log_dir, 'res/evaluation_pic', evaluation_dir ]) writer = SummaryWriter(log_dir) if os.path.exists(model_name) and args.load_model: print('Loading model ...') mol = torch.load(model_name).to(device) else: print('Init model ...') mol = ModelCls(args).to(device) ################################################################ # Prepare Data & Optimizer ################################################################ print('Prepare data loader ...') test_loader = getDataLoader(args, kwargs, train='test') train_loader = getDataLoader(args, kwargs, train='train') cover_val_loader = getDataLoader(args, kwargs, train='cover_validation') cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample') optimizer = torch.optim.Adam(mol.classifier.parameters(), lr=args.lr) loss_class = nn.CrossEntropyLoss(reduction='none').cuda(cuda) total, correct, top5correct, cnt = 0, 0, 0, 0 print('Start training ...') for epoch in range(args.epoch): ################################################################ # Testing ################################################################ if epoch > 0: mol.eval() # Testing on ImageNet val print('######### Testing on ImageNet val Dataset ###########') test_loss_cls, test_acc, test_top5acc = test_cls( test_loader, mol, cuda, 'Full') writer.add_scalar('test_imagenet/loss_classifier', test_loss_cls, epoch) writer.add_scalar('test_imagenet/accuracy', test_acc, epoch) writer.add_scalar('test_imagenet/top5accuracy', test_top5acc, epoch) if epoch % 5 == 4 or epoch == args.epoch - 1: mol.eval() evaluate_labeled_data(test_loader, mol, cuda) ################################################################ # Training ################################################################ step_time = time.time() mol.train() print('######### Training with %d batches total ##########' % len(train_loader)) for step, (x, y) in enumerate(train_loader): b_x = Variable(x).cuda() if cuda else Variable(x) label = Variable( torch.Tensor([y[2][i] for i in range(len(y[0]))]).long()) label = label.cuda() if cuda else label prob_class = mol(b_x) loss = loss_class(prob_class, label) writer.add_scalar('train/loss_classifier', torch.mean(loss.data), cnt) optimizer.zero_grad() torch.mean(loss).backward() optimizer.step() _, predicted = torch.max(prob_class.data, 1) total += label.size(0) correct += (predicted == label).sum().item() top5pre = prob_class.topk(5, 1, True, True) top5pre = top5pre[1].t() top5correct += top5pre.eq(label.view( 1, -1).expand_as(top5pre)).sum().item() writer.add_scalar('train/accuracy', correct / total, cnt) writer.add_scalar('train/top5_accuracy', top5correct / total, cnt) if step % 10 == 0: if os.path.exists(model_name): shutil.copy2(model_name, model_name.split('.pkl')[0] + '_back.pkl') torch.save(mol, model_name) total_tmp = total if total != 0 else 1 print( '[Training] Epoch:', epoch, 'Step:', step, '|', 'Time cost %.2f s; Classification error %.6f; ' 'Accuracy %.3f%%; Top5 Accuracy %.3f%%' % (time.time() - step_time, torch.mean(loss.data), correct * 100 / total_tmp, top5correct * 100 / total_tmp)) correct, total, top5correct = 0, 0, 0 step_time = time.time() cnt += 1 pass
def train(): ################################################################ # Arguments ################################################################ vae_args = train_args() cuda = vae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if vae_args.cuda else {} start_time = time.time() args = vae_args model_name = 'model/VAE_%s%s_model-%s.pkl' % ( args.model, '' if args.fea_c is None else args.fea_c, args.dataset) pic_dir = 'res/VAE_%s%s-%s/' % (args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') if cuda: vae = torch.load(model_name).to(device) else: vae = torch.load(model_name, map_location='cpu') else: vae = init_model(args.model).to(device) train_loader = getDataLoader(args, kwargs) optimizer = torch.optim.Adam(list(vae.parameters()), lr=args.lr) check_dir_exists(['res/', 'model', pic_dir]) loss_val = None for epoch in range(args.epoch): step_time = time.time() for step, (x, y) in enumerate(train_loader): b_x = Variable(x).cuda() if cuda else Variable(x) b_y = b_x.detach().cuda() if cuda else b_x.detach( ) # batch y, shape (batch, 32*32*3) decoded, mu, std = vae(b_x) if step % 100 == 0: img_to_save = torch.cat([b_x.data, decoded.data]) save_image(img_to_save, '%s/%s-%s.jpg' % (pic_dir, epoch, step)) # io.imsave('.xxx.jpg',img_to_save[0]) loss, bce, kld = loss_function(decoded, b_y, mu, std) # mean square error optimizer.zero_grad() # clear gradients for this training step loss.backward() # backpropagation, compute gradients optimizer.step() loss_val = 0.99 * loss_val + 0.01 * loss.data[ 0] if loss_val is not None else loss.data[0] if step % 10 == 0: torch.save(vae, model_name) print( 'Epoch:', epoch, 'Step:', step, '|', 'train loss %.6f; KLD %.6f; BCE %.6f; Time cost %.2f s' % (loss_val, kld, bce, time.time() - step_time)) step_time = time.time() print('Finished. Totally cost %.2f' % (time.time() - start_time))
def train(mol_short='inception_v3'): ################################################################ # Arguments ################################################################ ae_args = train_args() cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} # global ae_args, cuda, device, kwargs start_time = time.time() args = ae_args model_name = 'model/%s_%s%s_model-%s.pkl' % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % ( mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') mol = torch.load(model_name).to(device) else: print('Init model ...') mol = inception_v3(pretrained=True, training=False).to(device) # mol.eval() # train_loader = getDataLoader(args, kwargs) test_loader = getDataLoader(args, kwargs, train='test') # cover_loader = getDataLoader(args, kwargs, train='cover') cover_val_loader = getDataLoader(args, kwargs, train='cover_validation') cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample') loss_class = nn.CrossEntropyLoss().cuda(cuda) check_dir_exists(['res/', 'model', 'res/evaluation_pic', evaluation_dir]) # Evaluation # check_dir_exists([os.path.join(evaluation_dir, 'cos'), os.path.join(evaluation_dir, 'distance')]) evaluate_cover(cover_val_loader, cover_sample_loader, mol, cuda, evaluation_dir, args) # encode_accuracy, encode_top5accuracy, fc_accuracy, fc_top5accuracy = evaluate_labeled_data(test_loader, mol, cuda) # print('Encode accuracy:', np.mean(encode_accuracy)) # print('Encode top5 accuracy:', np.mean(encode_top5accuracy)) # print('Fc accuracy:', np.mean(fc_accuracy)) # print('Fc top5 accuracy:', np.mean(fc_top5accuracy)) total, correct, top5correct, loss_total = 0, 0, 0, 0 for epoch in range(1): step_time = time.time() for step, (x, y) in enumerate(test_loader): b_x = Variable(x, volatile=True).cuda() if cuda else Variable(x) label = Variable( torch.Tensor([y[2][i] for i in range(len(y[0]))]).long()) label = label.cuda() if cuda else label prob_class = mol(b_x) loss = loss_class(prob_class, label) # mean square error _, predicted = torch.max(prob_class.data, 1) total += label.size(0) correct += (predicted == label).sum().item() top5pre = prob_class.topk(5, 1, True, True) top5pre = top5pre[1].t() top5correct += top5pre.eq(label.view( 1, -1).expand_as(top5pre)).sum().item() loss_total += loss.data[0] * label.size(0) if step % 10 == 0: torch.save(mol, model_name) print( 'Epoch:', epoch, 'Step:', step, '|', 'test loss %.6f; Time cost %.2f s; Classification error %.6f; ' 'Top1 Accuracy %.3f; Top5 Accuracy %.3f' % (loss_total / total, time.time() - step_time, loss, correct * 100 / total, top5correct * 100 / total)) step_time = time.time() print('Finished. Totally cost %.2f' % (time.time() - start_time))
def train(mol_short='VGGClass', main_model=VGGClass): ################################################################ # Arguments ################################################################ ae_args = train_args() cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} # global ae_args, cuda, device, kwargs args = ae_args mol_short = mol_short if args.name == '' else mol_short + '_' + args.name log_dir = 'log/log_%s_%s%s_model-%s/' %\ (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) writer = SummaryWriter(log_dir) start_time = time.time() model_name = 'model/%s_%s%s_model-%s.pkl' % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) print('[Model] model name is', model_name) pic_dir = 'res/%s_%s%s-%s/' % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % ( mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') if cuda: mol = torch.load(model_name).to(device) else: mol = torch.load(model_name, map_location='cpu') else: print('Init model ...') mol = main_model(args.fea_c).to(device) print('Prepare data loader ...') train_loader = getDataLoader(args, kwargs, p=args.imgnet_p) test_loader = getDataLoader(args, kwargs, train='test') # small_test_loader = getDataLoader(args, kwargs, train=False, p=10) # cover_loader = getDataLoader(args, kwargs, train='cover') cover_val_loader = getDataLoader(args, kwargs, train='cover_validation') cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample') # Optimizer & Loss function if args.fea_c == 512: optimizer1 = torch.optim.Adam(list(mol.classification.parameters()), lr=args.lr) else: optimizer1 = torch.optim.Adam(list(mol.classification.parameters()) + list(mol.small_features.parameters()), lr=args.lr) loss_class = nn.CrossEntropyLoss().cuda(cuda) loss_val = None # Check directories check_dir_exists([ 'res/', 'model', pic_dir, log_dir, 'res/evaluation_pic', evaluation_dir ]) total, correct, top5correct, cnt = 0, 0, 0, 0 print('Start training ...') for epoch in range(args.epoch): # Evaluation cover if epoch % 5 == 0 and epoch != 0: mol.eval() eval_dir = os.path.join(evaluation_dir, 'epoch%d' % epoch) evaluate_cover(cover_val_loader, cover_sample_loader, mol, cuda, eval_dir, args) encode_accuracy, encode_top5accuracy, fc_accuracy, fc_top5accuracy = evaluate_labeled_data( test_loader, mol, cuda) writer.add_scalar('test/encode_feature_accuracy', np.mean(encode_accuracy), epoch) writer.add_scalar('test/encode_feature_top5accuracy', np.mean(encode_top5accuracy), epoch) writer.add_scalar('test/fc_feature_accuracy', np.mean(fc_accuracy), epoch) writer.add_scalar('test/fc_feature_top5accuracy', np.mean(fc_top5accuracy), epoch) if epoch != 0: # Testing classifier mol.eval() test_acc, test_top5acc = test(test_loader, mol, cuda, 'Full') writer.add_scalar('test/class_accuracy', test_acc, epoch) writer.add_scalar('test/class_top5accuracy', test_top5acc, epoch) step_time = time.time() mol.train() print('######### Training with %d batches total ##########' % len(train_loader)) for step, (x, y) in enumerate(train_loader): b_x = Variable(x).cuda() if cuda else Variable(x) label = Variable( torch.Tensor([y[2][i] for i in range(len(y[0]))]).long()) label = label.cuda() if cuda else label prob_class = mol(b_x) loss = loss_class(prob_class, label) writer.add_scalar('train/loss_classifier', loss, cnt) optimizer1.zero_grad() loss.backward() optimizer1.step() _, predicted = torch.max(prob_class.data, 1) total += label.size(0) correct += (predicted == label).sum().item() top5pre = prob_class.topk(5, 1, True, True) top5pre = top5pre[1].t() top5correct += top5pre.eq(label.view( 1, -1).expand_as(top5pre)).sum().item() writer.add_scalar('train/accuracy', correct / total, cnt) writer.add_scalar('train/top5_accuracy', top5correct / total, cnt) loss_val = 0.99 * loss_val + 0.01 * loss.data[ 0] if loss_val is not None else loss.data[0] if step % 10 == 0: if os.path.exists(model_name): shutil.copy2(model_name, model_name.split('.pkl')[0] + '_back.pkl') torch.save(mol, model_name) print( '[Training] Epoch:', epoch, 'Step:', step, '|', 'train loss %.6f; Time cost %.2f s; Accuracy %.3f%%; Top5 Accuracy %.3f%%' % (loss.data[0], time.time() - step_time, correct * 100 / total, top5correct * 100 / total)) correct, total, top5correct = 0, 0, 0 step_time = time.time() cnt += 1 print('Finished. Totally cost %.2f' % (time.time() - start_time)) writer.export_scalars_to_json(os.path.join(log_dir, 'all_scalars.json')) writer.close()
def train(): ################################################################ # Arguments ################################################################ ae_args = train_args() cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} # global ae_args, cuda, device, kwargs start_time = time.time() args = ae_args model_name = 'model/vgg_classifier_%s%s_model-%s.pkl' % ( args.model, '' if args.fea_c is None else args.fea_c, args.dataset) evaluation_dir = 'res/evaluation_pic/vgg_%s%s-%s' % ( args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') vgg = torch.load(model_name).to(device) else: print('Init model ...') vgg = VGGNet(args.fea_c).to(device) # train_loader = getDataLoader(args, kwargs) test_loader = getDataLoader(args, kwargs, train='test') # cover_loader = getDataLoader(args, kwargs, train='cover') cover_val_loader = getDataLoader(args, kwargs, train='cover_validation') cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample') optimizer = torch.optim.Adam(list(vgg.parameters()), lr=args.lr) loss_class = nn.CrossEntropyLoss().cuda(cuda) check_dir_exists(['res/', 'model', 'res/evaluation_pic', evaluation_dir]) # # Evaluation # vgg.eval() # check_dir_exists([os.path.join(evaluation_dir, 'cos'), os.path.join(evaluation_dir, 'distance')]) # # evaluate_cover(cover_val_loader, cover_sample_loader, vgg, cuda, evaluation_dir, args) # fc_accuracy, fc_top5accuracy = evaluate_labeled_data(test_loader, vgg, cuda, both=False) # # print('Encode accuracy:', np.mean(encode_accuracy)) # # print('Encode top5 accuracy:', np.mean(encode_top5accuracy)) # print('Fc accuracy:', np.mean(fc_accuracy)) # print('Fc top5 accuracy:', np.mean(fc_top5accuracy)) total, correct, top5correct, loss_total = 0, 0, 0, 0 t_per_img = [] for epoch in range(1): step_time = time.time() for step, (x, y) in enumerate(test_loader): b_x = Variable(x, volatile=True).cuda() if cuda else Variable(x) label = Variable( torch.Tensor([y[2][i] for i in range(len(y[0]))]).long()) label = label.cuda() if cuda else label t0 = time.time() prob_class = vgg(b_x) t_tmp = (time.time() - t0) / len(b_x) * 1000 t_per_img.append(t_tmp) print( 'cost %.6fms per image this batch. cost %.6fms per image till now.' % (t_tmp, np.mean(sorted(t_per_img)[1:-1]))) loss = loss_class(prob_class, label) # mean square error # optimizer.zero_grad() # clear gradients for this training step # loss.backward() # optimizer.step() _, predicted = torch.max(prob_class.data, 1) total += label.size(0) correct += (predicted == label).sum().item() top5pre = prob_class.topk(5, 1, True, True) top5pre = top5pre[1].t() top5correct += top5pre.eq(label.view( 1, -1).expand_as(top5pre)).sum().item() loss_total += loss.data[0] * label.size(0) if step % 10 == 0: torch.save(vgg, model_name) print( 'Epoch:', epoch, 'Step:', step, '|', 'test loss %.6f; Time cost %.2f s; Classification error %.6f; ' 'Top1 Accuracy %.3f; Top5 Accuracy %.3f' % (loss_total / total, time.time() - step_time, loss.item(), correct * 100 / total, top5correct * 100 / total)) step_time = time.time() loss = None print('Finished. Totally cost %.2f' % (time.time() - start_time))
def train(): ################################################################ # Arguments ################################################################ ae_args = args = feature_classifier_args() cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} ################################################################ # Load or Init model ################################################################ start_time = time.time() model_name = 'model/FeaClass_%s_%s%s_model-%s.pkl' \ % (args.main_model, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') mol = torch.load(model_name).to(device) else: print('Init model ...') feature_name = 'model/%s_%s%s_model-%s.pkl' % ( args.main_model, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) print('Loading feature model (%s)...' % feature_name) assert os.path.exists(feature_name) if cuda: fea_mol = torch.load(feature_name).to(device) else: fea_mol = torch.load(feature_name, map_location='cpu').to(device) mol = FClassifier(fea_mol, args.main_model, args.fea_c).to(device) train_loader = getDataLoader(args, kwargs) optimizer = torch.optim.Adam(list(mol.classification.parameters()), lr=args.lr) loss_fn = nn.CrossEntropyLoss().cuda(cuda) check_dir_exists(['res/', 'model']) total, correct, top5correct = 0, 0, 0 for epoch in range(args.epoch): step_time = time.time() for step, (x, y) in enumerate(train_loader): b_x = Variable(x).cuda() if cuda else Variable(x) label = Variable( torch.Tensor([y[2][i] for i in range(len(y[0]))]).long()) label = label.cuda() if cuda else label prob_class = mol(b_x) loss = loss_fn(prob_class, label) optimizer.zero_grad() loss.backward() optimizer.step() _, predicted = torch.max(prob_class.data, 1) total += label.size(0) correct += (predicted == label).sum().item() top5pre = prob_class.topk(5, 1, True, True) top5pre = top5pre[1].t() top5correct += top5pre.eq(label.view( 1, -1).expand_as(top5pre)).sum().item() if step % 10 == 0: shutil.copy2(model_name, model_name.split('.pkl')[0] + '_back.pkl') torch.save(mol, model_name) print( 'Epoch:', epoch, 'Step:', step, '|', 'train loss %.6f; Time cost %.2f s; Classification error %.6f' 'Accuracy %.3f%%; Top5 Accuracy %.3f%%' % (loss.data[0], time.time() - step_time, loss, correct * 100 / total, top5correct * 100 / total)) correct, total, top5correct = 0, 0, 0 step_time = time.time() print('Finished. Totally cost %.2f' % (time.time() - start_time))
def train(args, mol_short='AEClass_both', main_model=AEClass): print('######### This is train function #############') ################################################################ # Arguments ################################################################ ae_args = args if args.gpu == -1: os.environ["CUDA_VISIBLE_DEVICES"] = GPU().choose_gpu() elif args.gpu is not None: os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} d_name = '' if args.decoder == 'vgg' else '_' + args.decoder + 'decoder' mol_short = mol_short + args.name + d_name log_dir = 'log/log_%s_%s%s_model-%s/' %\ (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) remove_dir_exists([log_dir]) writer = SummaryWriter(log_dir) start_time = time.time() model_name = 'model/%s_%s%s_model-%s.pkl' % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) print('[Model] model name is', model_name) pic_dir = 'res/%s_%s%s-%s/' % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % ( mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') if cuda: mol = torch.load(model_name).to(device) else: mol = torch.load(model_name, map_location='cpu') else: print('Init model ...') mol = main_model(args.fea_c).to(device) print('Prepare data loader ...') # train_loader = getDataLoader(args, kwargs, train='train') fuse_loader = getDataLoader(args, kwargs, train='fuse', p=args.imgnet_p) test_dataset = RandomBatchSampler(getDataset(args, train='test'), args.batch_size) test_loader = getDataLoader(args, kwargs, train='test') # small_test_loader = getDataLoader(args, kwargs, train=False, p=10) # cover_loader = getDataLoader(args, kwargs, train='cover') cover_val_loader = getDataLoader(args, kwargs, train='cover_validation') cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample') # optimizer_cls = torch.optim.Adam(get_optimized_params(mol, 'classifier', lr=args.lr) + # get_optimized_params(mol, 'small_features', lr=args.lr) + # get_optimized_params(mol, 'vgg.38', lr=args.lr/5), lr=args.lr) optimizer_d = torch.optim.Adam( ([] if args.fea_c == 512 else get_optimized_params( mol, 'small_features', args.lr)) + get_optimized_params(mol, 'decoder', args.lr) + get_optimized_params(mol, 'vgg.38', args.lr / 5), lr=args.lr) optimizer = torch.optim.Adam( ([] if args.fea_c == 512 else get_optimized_params( mol, 'small_features', args.lr)) + get_optimized_params(mol, 'classifier', lr=args.lr) + get_optimized_params(mol, 'decoder', args.lr) + get_optimized_params(mol, 'vgg.38', args.lr / 5), lr=args.lr) loss_decoder = nn.MSELoss() loss_class = nn.CrossEntropyLoss(reduction='none').cuda(cuda) def loss_cls_fn(prob_class, label, weight, beta=0.01): if torch.sum(weights) == 0: return 99999 else: lss = loss_class(prob_class, label) * weight lss = torch.sum(lss) / torch.sum(weight) return lss * beta check_dir_exists([ 'res/', 'model', pic_dir, log_dir, 'res/evaluation_pic', evaluation_dir ]) total, correct, top5correct, cnt = 0, 0, 0, 0 print('Start training ...') for epoch in range(args.epoch): if (epoch % 5 == 0) and epoch != 0: # Evaluation on cover data mol.eval() eval_dir = os.path.join(evaluation_dir, 'epoch%d' % epoch) evaluate_cover(cover_val_loader, cover_sample_loader, mol, cuda, eval_dir, args) fc_accuracy, fc_top5accuracy = evaluate_labeled_data(test_loader, mol, cuda, both=False) print('Fc accuracy:', np.mean(fc_accuracy)) print('Fc top5 accuracy:', np.mean(fc_top5accuracy)) if epoch >= 0: mol.eval() # Testing on ImageNet val print('######### Testing on ImageNet val Dataset ###########') test_loss_decoder, test_loss_cls, test_acc, test_top5acc = test_cls_decoder( test_loader, mol, cuda, 'Full') # test_loss = (1 - args.alpha) * test_loss_cls + args.alpha * test_loss_decoder / 0.001 writer.add_scalar('test_imagenet/loss_decoder', test_loss_decoder, epoch) writer.add_scalar('test_imagenet/loss_classifier', test_loss_cls, epoch) # writer.add_scalar('test_imagenet/loss', test_loss, epoch) writer.add_scalar('test_imagenet/accuracy', test_acc, epoch) writer.add_scalar('test_imagenet/top5accuracy', test_top5acc, epoch) # Testing on Cover val print('######### Testing on Cover val Dataset ###########') # test_loss_decoder, test_loss_cls, test_acc, test_top5acc = test_cls_decoder(cover_val_loader, mol, cuda, 'Full') test_loss_decoder = test_decoder(cover_val_loader, mol, cuda, 'Full') # test_loss = (1 - args.alpha) * test_loss_cls + args.alpha * test_loss_decoder / 0.001 writer.add_scalar('test_cover/loss_decoder', test_loss_decoder, epoch) # writer.add_scalar('test_cover/loss_classifier', test_loss_cls, epoch) # writer.add_scalar('test_cover/loss', test_loss, epoch) # writer.add_scalar('test_cover/accuracy', test_acc, epoch) # writer.add_scalar('test_cover/top5accuracy', test_top5acc, epoch) step_time = time.time() for step, (x, y) in enumerate(fuse_loader): mol.train() b_x = Variable(x).cuda() if cuda else Variable(x) b_y = b_x.detach().cuda() if cuda else b_x.detach( ) # batch y, shape (batch, 32*32*3) # //TODO # train_cls = label = Variable( torch.Tensor([y[2][i] for i in range(len(y[0]))]).long()) label = label.cuda() if cuda else label weights = Variable( torch.Tensor([y[3][i] for i in range(len(y[0]))])) weights = weights.cuda() if cuda else weights _, decoded, prob_class = mol(b_x) if step % 100 == 0: img_to_save = decoded.data save_image(img_to_save, '%s/%s-%s.jpg' % (pic_dir, epoch, step)) loss1 = loss_decoder(decoded, b_y) loss2 = loss_cls_fn(prob_class, label, weights, beta=0.01) loss = 0 writer.add_scalar('train/loss_decoder', loss1, cnt) writer.add_scalar('train/loss_classifier', loss2, cnt) # writer.add_scalar('train/loss', loss, cnt) if loss2 != 99999: loss = (1 - args.alpha) * loss2 + args.alpha * loss1 optimizer.zero_grad() loss.backward() optimizer.step() _, predicted = torch.max(prob_class.data, 1) predicted = torch.Tensor([ predicted[i] for i in range(len(weights)) if weights[i] == 1 ]) label2 = torch.Tensor( [label[i] for i in range(len(weights)) if weights[i] == 1]) total += predicted.shape[0] correct += (predicted == label2).sum().item() prob_class2 = prob_class[[ i for i in range(len(weights)) if weights[i] == 1 ]] top5pre = prob_class2.topk(5, 1, True, True) top5pre = top5pre[1].t() label2 = label2.long().cuda() if cuda else label2.long() top5correct += top5pre.eq( label2.view(1, -1).expand_as(top5pre)).sum().item() writer.add_scalar('train/accuracy', correct / total, cnt) writer.add_scalar('train/top5_accuracy', top5correct / total, cnt) else: optimizer_d.zero_grad() loss1.backward() optimizer_d.step() # loss_val = 0.99*loss_val + 0.01*loss.data[0] if loss_val is not None else loss.data[0] if step % 10 == 0: if os.path.exists(model_name): shutil.copy2(model_name, model_name.split('.pkl')[0] + '_back.pkl') torch.save(mol, model_name) total_tmp = total if total != 0 else 1 t_x, t_y = test_dataset.get_sample() test_acc, test_top5_acc = _test_sample_batch( mol, cuda, t_x, t_y) print( '[Training] Epoch:', epoch, 'Step:', step, '|', 'Time cost %.2f s; Classification error %.6f; Decoder error %.6f; Loss %.6f; ' 'Accuracy %.3f%%/%.3f%%; Top5 Accuracy %.3f%%/%.3f%% (%s)' % (time.time() - step_time, 0.0888 if type(loss2) == int else loss2.data[0], loss1.data[0], 0.0888 if type(loss) == int else loss.data[0], correct * 100 / total_tmp, test_acc * 100, top5correct * 100 / total_tmp, test_top5_acc * 100, time.ctime())) writer.add_scalar('test/accuracy', test_acc, cnt) writer.add_scalar('test/top5_accuracy', test_top5_acc, cnt) correct, total, top5correct = 0, 0, 0 step_time = time.time() cnt += 1 print('Finished. Totally cost %.2f' % (time.time() - start_time)) writer.export_scalars_to_json(os.path.join(log_dir, 'all_scalars.json')) writer.close()
def train_decoder_only(args, mol_short='AEClass_d', main_model=AEClass): print('######### This is train decoder only function #############') ################################################################ # Arguments ################################################################ ae_args = args cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} # global ae_args, cuda, device, kwargs d_name = '' if args.decoder == 'vgg' else '_' + args.decoder + 'decoder' mol_short = mol_short + args.name + d_name log_dir = 'log/log_%s_%s%s_model-%s/' %\ (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) remove_dir_exists([log_dir]) writer = SummaryWriter(log_dir) start_time = time.time() model_name = 'model/%s_%s%s_model-%s.pkl' \ % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) print('[Model] model name is', model_name) pic_dir = 'res/%s_%s%s-%s/' % (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % ( mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') mol = torch.load(model_name).to(device) else: print('Init model ...') mol = main_model(args.fea_c, decoder=args.decoder).to(device) print('Prepare data loader ...') test_loader = getDataLoader(args, kwargs, train='test') # small_test_loader = getDataLoader(args, kwargs, train=False, p=10) cover_loader = getDataLoader(args, kwargs, train='cover') cover_val_loader = getDataLoader(args, kwargs, train='cover_validation') cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample') # optimizer1 = torch.optim.Adam(list(mol.classification.parameters())+list(mol.small_features.parameters())+ # list(mol.decoder.parameters()), lr=args.lr) optimizer1 = torch.optim.Adam( ([] if args.fea_c == 512 else get_optimized_params( mol, 'small_features', args.lr)) + get_optimized_params(mol, 'decoder', args.lr) + get_optimized_params(mol, 'vgg.38', args.lr / 5), lr=args.lr) loss_decoder = nn.MSELoss() check_dir_exists([ 'res/', 'model', pic_dir, log_dir, 'res/evaluation_pic', evaluation_dir ]) loss_val = None print('Start training ...') cnt = 0 for epoch in range(args.epoch): if epoch % 5 == (5 - 1) or epoch == (args.epoch - 1): # Evaluation on cover data eval_dir = os.path.join(evaluation_dir, 'epoch%d' % epoch) evaluate_cover(cover_val_loader, cover_sample_loader, mol, cuda, eval_dir) if epoch != 0: # Testing on Cover val print('######### Testing on Cover val Dataset ###########') test_loss_decoder = test_decoder(cover_val_loader, mol, cuda, 'Full') # test_loss = (1 - args.alpha) * test_loss_cls + args.alpha * test_loss_decoder / 0.001 writer.add_scalar('test_cover/loss_decoder', test_loss_decoder, epoch) # writer.add_scalar('test_cover/loss_classifier', test_loss_cls, epoch) # writer.add_scalar('test_cover/loss', test_loss, epoch) # writer.add_scalar('test_cover/accuracy', test_acc, epoch) # writer.add_scalar('test_cover/top5accuracy', test_top5acc, epoch) # Testing on ImageNet val print('######### Testing on ImageNet val Dataset ###########') # test_loss_decoder, test_loss_cls, test_acc, test_top5acc = test_cls_decoder(test_loader, mol, cuda, 'Full') test_loss_decoder = test_decoder(test_loader, mol, cuda, 'Full') # test_loss = (1 - args.alpha) * test_loss_cls + args.alpha * test_loss_decoder / 0.001 writer.add_scalar('test_imagenet/loss_decoder', test_loss_decoder, epoch) # writer.add_scalar('test_imagenet/loss_classifier', test_loss_cls, epoch) # writer.add_scalar('test_imagenet/loss', test_loss, epoch) # writer.add_scalar('test_imagenet/accuracy', test_acc, epoch) # writer.add_scalar('test_imagenet/top5accuracy', test_top5acc, epoch) step_time = time.time() print('######### Training with %d batches total ##########' % len(cover_loader)) for step, (x, y) in enumerate(cover_loader): b_x = Variable(x).cuda() if cuda else Variable(x) b_y = b_x.detach().cuda() if cuda else b_x.detach( ) # batch y, shape (batch, 32*32*3) _, decoded, prob_class = mol(b_x) if step % 500 == 0 or step == (len(cover_loader) - 1): img_to_save = decoded.data save_image(img_to_save, '%s/%s-%s.jpg' % (pic_dir, epoch, step)) loss1 = loss_decoder(decoded, b_y) # loss = (1-args.alpha) * loss2 + args.alpha * loss1 / 0.001 loss = loss1 writer.add_scalar('train/loss_decoder', loss1, cnt) writer.add_scalar('train/loss', loss, cnt) optimizer1.zero_grad() loss.backward() optimizer1.step() loss_val = 0.99 * loss_val + 0.01 * loss.data[ 0] if loss_val is not None else loss.data[0] if step % 10 == 0 or step == (len(cover_loader) - 1): if os.path.exists(model_name): shutil.copy2(model_name, model_name.split('.pkl')[0] + '_back.pkl') torch.save(mol, model_name) print( '[Training] Epoch:', epoch, 'Step:', step, '|', 'train loss %.6f; Time cost %.2f s; Decoder error %.6f;' % (loss.data[0], time.time() - step_time, loss1.data[0])) step_time = time.time() cnt += 1 print('Finished. Totally cost %.2f' % (time.time() - start_time)) writer.export_scalars_to_json(os.path.join(log_dir, 'all_scalars.json')) writer.close()
def train(): ################################################################ # Arguments ################################################################ ae_args = train_args() cuda = ae_args.cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {} # global ae_args, cuda, device, kwargs args = ae_args log_dir = 'log/log_AE_%s%s_model-%s/' %\ (args.model, '' if args.fea_c is None else args.fea_c, args.dataset) writer = SummaryWriter(log_dir) start_time = time.time() model_name = 'model/AE_%s%s_model-%s.pkl' % ( args.model, '' if args.fea_c is None else args.fea_c, args.dataset) pic_dir = 'res/AE_%s%s-%s/' % (args.model, '' if args.fea_c is None else args.fea_c, args.dataset) if os.path.exists(model_name) and args.load_model: print('Loading model ...') mol = torch.load(model_name).to(device) else: print('Init model ...') mol = AutoEncoder(args.fea_c).to(device) print('Prepare data loader ...') train_loader = getDataLoader(args, kwargs, train=True) test_loader = getDataLoader(args, kwargs, train=False) optimizer1 = torch.optim.Adam(list(mol.small_features.parameters()) + list(mol.decoder.parameters()), lr=args.lr) optimizer2 = torch.optim.Adam(list(mol.features.parameters()), lr=args.lr / 10) loss_decoder = nn.MSELoss() check_dir_exists(['res/', 'model', pic_dir, log_dir]) total, correct, top5correct, cnt = 0, 0, 0, 0 print('Start training ...') for epoch in range(args.epoch): # Testing test_acc, test_top5acc, test_loss = test_feature( test_loader, mol, cuda, 'Full') writer.add_scalar('test/accuracy', np.mean(test_acc), epoch) writer.add_scalar('test/top5accuracy', np.mean(test_top5acc), epoch) writer.add_scalar('test/loss_decoder', test_loss, epoch) step_time = time.time() for step, (x, y) in enumerate(train_loader): b_x = Variable(x).cuda() if cuda else Variable(x) b_y = b_x.detach().cuda() if cuda else b_x.detach( ) # batch y, shape (batch, 32*32*3) encoded, decoded = mol(b_x) if step % 100 == 0: img_to_save = decoded.data save_image(img_to_save, '%s/%s-%s.jpg' % (pic_dir, epoch, step)) loss = loss_decoder(decoded, b_y) writer.add_scalar('train/loss_decoder', loss, cnt) optimizer1.zero_grad() optimizer2.zero_grad() loss.backward() optimizer1.step() optimizer2.step() if step % 50 == 0: if os.path.exists(model_name): shutil.copy2(model_name, model_name.split('.pkl')[0] + '_back.pkl') torch.save(mol, model_name) print( '[Training] Epoch:', epoch, 'Step:', step, '|', 'train loss %.6f; Time cost %.2f s; Decoder error %.6f' % (loss.data[0], time.time() - step_time, loss)) step_time = time.time() cnt += 1 print('Finished. Totally cost %.2f' % (time.time() - start_time)) writer.export_scalars_to_json(os.path.join(log_dir, 'all_scalars.json')) writer.close()