def main_wflwe70(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) dataset_type = get_dataset(config) dataset = dataset_type(config, is_train=True) for i in range(len(dataset)): # ipdb.set_trace() img, fname, meta = dataset[i] filename = osp.join('data/wflwe70/xximages', fname) if not osp.exists(osp.dirname(filename)): os.makedirs(osp.dirname(filename)) scale = meta['scale'] center = meta['center'] tpts = meta['tpts'] for spt in tpts: img = cv2.circle(img, (4 * spt[0], 4 * spt[1]), 1 + center[0] // 400, (255, 0, 0)) img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite(filename, img)
def main_cofw(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) ipdb.set_trace() dataset_type = get_dataset(config) dataset = dataset_type(config, is_train=True) fp = open('data/cofw/test.csv', 'w') for i in range(len(dataset)): # ipdb.set_trace() img, image_path, meta = dataset[i] fname = osp.join('data/cofw/test', osp.basename(image_path)) fp.write('%s,1,128,128' % fname) tpts = meta['tpts'] for j in range(tpts.shape[0]): fp.write(',%d,%d' % (tpts[j, 0], tpts[j, 1])) fp.write('\n') img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite(fname, img) fp.close()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) gpus = list(config.GPUS) model = nn.DataParallel(model, device_ids=gpus).cuda() # load model state_dict = torch.load(args.model_file) if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict'] model.load_state_dict(state_dict) else: model.module.load_state_dict(state_dict) dataset_type = get_dataset(config) test_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) nme, predictions = function.inference(config, test_loader, model) with open('../data/crabs/crabs_data_test.csv', 'r') as f: data = np.loadtxt(f, str, delimiter=",", skiprows=1) paths = data[:, 0] for index, path in enumerate(paths): img = cv2.imread("../data/crabs/images/{}".format(path)) a = predictions[index] b = a.numpy() for index, px in enumerate(b): # print(tuple(px)) cv2.circle(img, tuple(px), 1, (0, 0, 255), 3, 8, 0) # cv2.imwrite("/home/njtech/Jiannan/crabs/dataset/result_new/{}".format(path.split('/')[-1]), img) cv2.imshow("img", img) cv2.waitKey(1000) & 0xFF torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) if args.onnx_export: torch_out = torch.onnx._export(model, torch.rand(1, 3, config.IMAGE_SIZE), osp.join(final_output_dir, args.onnx_export), export_params=True) return gpus = list(config.GPUS) if gpus[0] > -1: model = nn.DataParallel(model, device_ids=gpus).cuda() # load model if gpus[0] > -1: state_dict = torch.load(args.model_file) else: state_dict = torch.load(args.model_file, map_location='cpu') if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict'] model.load_state_dict(state_dict) else: if gpus[0] > -1: model.module.load_state_dict(state_dict) else: model.load_state_dict(state_dict) dataset_type = get_dataset(config) dataset = dataset_type(config, is_train=False) test_loader = DataLoader(dataset=dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) ipdb.set_trace() nme, predictions = function.inference(config, test_loader, model) torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() # model = models.get_face_alignment_net(config) model = eval('models.' + config.MODEL.NAME + '.get_face_alignment_net')( config, is_train=True) gpus = list(config.GPUS) model = nn.DataParallel(model, device_ids=gpus).cuda() # load model # state_dict = torch.load(args.model_file) # if 'state_dict' in state_dict.keys(): # state_dict = state_dict['state_dict'] # model.load_state_dict(state_dict) # else: # model.module.load_state_dict(state_dict) if args.model_file: logger.info('=> loading model from {}'.format(args.model_file)) # model.load_state_dict(torch.load(args.model_file), strict=False) model_state = torch.load(args.model_file) model.module.load_state_dict(model_state.state_dict()) else: model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info('=> loading model from {}'.format(model_state_file)) model_state = torch.load(model_state_file) model.module.load_state_dict(model_state) dataset_type = get_dataset(config) test_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) nme, predictions = function.inference(config, test_loader, model) torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
def eval_text(self): self.set_mode('eval') tst_loader = DataLoader(get_dataset('iam_word', self.opt.training.dset_split), batch_size=self.opt.test.nrow, shuffle=True, collate_fn=self.collect_fn) def get_space_index(text): idxs = [] for i, ch in enumerate(text): if ch == ' ': idxs.append(i) return idxs with torch.no_grad(): nrow = self.opt.test.nrow while True: text = input('input text: ') if len(text) == 0: break batch = next(iter(tst_loader)) imgs, img_lens, lbs, lb_lens, wids = batch real_imgs, real_img_lens = imgs.to(self.device), img_lens.to( self.device) fake_lbs = self.label_converter.encode(text) fake_lbs = torch.LongTensor(fake_lbs) fake_lb_lens = torch.IntTensor([len(text)]) fake_lbs = fake_lbs.repeat(nrow, 1).to(self.device) fake_lb_lens = fake_lb_lens.repeat(nrow, ).to(self.device) enc_styles = self.models.E(real_imgs, real_img_lens, self.models.W.cnn_backbone) noises = torch.randn((nrow, self.noise_dim)).to(self.device) enc_styles = torch.cat([noises, enc_styles], dim=-1) real_imgs = (1 - real_imgs).squeeze().cpu().numpy() * 127 gen_imgs = self.models.G(enc_styles, fake_lbs, fake_lb_lens) gen_imgs = (1 - gen_imgs).squeeze().cpu().numpy() * 127 space_indexs = get_space_index(text) for idx in space_indexs: gen_imgs[:, :, idx * 16:(idx + 1) * 16] = 255 plt.figure() for i in range(nrow): plt.subplot(nrow * 2, 1, i * 2 + 1) plt.imshow(real_imgs[i], cmap='gray') plt.axis('off') plt.subplot(nrow * 2, 1, i * 2 + 2) plt.imshow(gen_imgs[i], cmap='gray') plt.axis('off') plt.tight_layout() plt.show()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) gpus = list(config.GPUS) model = nn.DataParallel(model, device_ids=gpus).cuda() # load model #state_dict = torch.load(args.model_file) #if 'state_dict' in state_dict.keys(): # state_dict = state_dict['state_dict'] # model.load_state_dict(state_dict) #else: # model.module.load_state_dict(state_dict) model = torch.load(args.model_file) model.eval() model = nn.DataParallel(model, device_ids=gpus).cuda() dataset_type = get_dataset(config) test_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) nme, predictions = function.inference(config, test_loader, model) import cv2 img = cv2.imread('data/wflw/images/my3.jpg') print(predictions, predictions.shape) for item in predictions[0]: cv2.circle(img, (item[0], item[1]), 3, (0, 0, 255), -1) cv2.imwrite('out.png', img) torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
def validate(self, guided=True): self.set_mode('eval') dset_name = self.opt.valid.dset_name if self.opt.valid.dset_name \ else self.opt.dataset dset = get_dataset(dset_name, self.opt.valid.dset_split) dloader = DataLoader(dset, collate_fn=self.collect_fn, batch_size=self.opt.valid.batch_size, shuffle=False, num_workers=4) # style images are resized source_dloader = DataLoader(get_dataset( self.opt.valid.dset_name.strip('_org'), self.opt.valid.dset_split), collate_fn=self.collect_fn, batch_size=self.opt.valid.batch_size, shuffle=False, num_workers=4) generator = self.image_generator(source_dloader, guided) fid_kid = calculate_kid_fid(self.opt.valid, dloader, generator, self.max_valid_image_width, self.device) return fid_kid
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) gpus = list(config.GPUS) if torch.cuda.is_available(): model = nn.DataParallel(model, device_ids=gpus).cuda() # load model state_dict = torch.load(args.model_file) else: # model = nn.DataParallel(model) state_dict = torch.load(args.model_file, map_location=lambda storage, loc: storage) if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict'] model.load_state_dict(state_dict) else: try: model.module.load_state_dict(state_dict.state_dict()) except AttributeError: state_dict ## remove first seven model.load_state_dict(state_dict) dataset_type = get_dataset(config) test_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) #nme, predictions = function.inference(config, test_loader, model, args.model_file) #### testing function.test(config, test_loader, model, args.model_file) #### testing
def eval_style(self): self.set_mode('eval') tst_loader = DataLoader(get_dataset('iam_word', self.opt.training.dset_split), batch_size=self.opt.test.nrow, shuffle=True, collate_fn=self.collect_fn) with torch.no_grad(): nrow, ncol = self.opt.test.nrow, 2 while True: text = input('input text: ') if len(text) == 0: break texts = text.split(' ') ncol = len(texts) batch = next(iter(tst_loader)) imgs, img_lens, lbs, lb_lens, wids = batch real_imgs, real_img_lens = imgs.to(self.device), img_lens.to( self.device) if len(texts) == 1: fake_lbs = self.label_converter.encode(texts) fake_lbs = torch.LongTensor(fake_lbs) fake_lb_lens = torch.IntTensor([len(texts[0])]) else: fake_lbs, fake_lb_lens = self.label_converter.encode(texts) fake_lbs = fake_lbs.repeat(nrow, 1).to(self.device) fake_lb_lens = fake_lb_lens.repeat(nrow, ).to(self.device) enc_styles = self.models.E(real_imgs, real_img_lens, self.models.W.cnn_backbone).unsqueeze(1).\ repeat(1, ncol, 1).view(nrow * ncol, self.opt.EncModel.style_dim) noises = torch.randn((nrow, self.noise_dim)).unsqueeze(1).\ repeat(1, ncol, 1).view(nrow * ncol, self.noise_dim).to(self.device) enc_styles = torch.cat([noises, enc_styles], dim=-1) gen_imgs = self.models.G(enc_styles, fake_lbs, fake_lb_lens) gen_imgs = (1 - gen_imgs).squeeze().cpu().numpy() * 127 real_imgs = (1 - real_imgs).squeeze().cpu().numpy() * 127 plt.figure() for i in range(nrow): plt.subplot(nrow, 1 + ncol, i * (1 + ncol) + 1) plt.imshow(real_imgs[i], cmap='gray') plt.axis('off') for j in range(ncol): plt.subplot(nrow, 1 + ncol, i * (1 + ncol) + 2 + j) plt.imshow(gen_imgs[i * ncol + j], cmap='gray') plt.axis('off') plt.tight_layout() plt.show()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = torchvision.models.resnet101(pretrained=config.MODEL.PRETRAINED, progress=True) num_ftrs = model.fc.in_features model.fc = torch.nn.Linear(num_ftrs, config.MODEL.OUTPUT_SIZE[0]) gpus = list(config.GPUS) model = nn.DataParallel(model, device_ids=gpus).cuda() # load model state_dict = torch.load(args.model_file) if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict'] model.load_state_dict(state_dict) else: model.module.load_state_dict(state_dict) dataset_type = get_dataset(config) test_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) predictions = function.inference(config, test_loader, model) torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
def main_300w(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) dataset_type = get_dataset(config) dataset = dataset_type(config, is_train=False) fp = open('data/300w/face_landmarks70_300w_test.csv', 'w') for i in range(len(dataset)): # ipdb.set_trace() img, fname, meta = dataset[i] filename = osp.join('data/300w/xximages', fname) if not osp.exists(osp.dirname(filename)): os.makedirs(osp.dirname(filename)) scale = meta['scale'] center = meta['center'] tpts = meta['tpts'] selpts = [] for j in range(0, 68): selpts.append(tpts[j]) selpts.append(tpts[36:42].mean(0)) selpts.append(tpts[42:48].mean(0)) fp.write('%s,%.2f,%.1f,%.1f' % (fname, scale, center[0], center[1])) for spt in selpts: img = cv2.circle(img, (spt[0], spt[1]), 1 + center[0] // 400, (255, 0, 0)) fp.write(',%f,%f' % (spt[0], spt[1])) fp.write('\n') img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite(filename, img) fp.close()
def main(hparams): results_dir = get_results_directory(hparams.output_dir) writer = SummaryWriter(log_dir=str(results_dir)) ds = get_dataset(hparams.dataset, root=hparams.data_root) input_size, num_classes, train_dataset, test_dataset = ds hparams.seed = set_seed(hparams.seed) if hparams.n_inducing_points is None: hparams.n_inducing_points = num_classes print(f"Training with {hparams}") hparams.save(results_dir / "hparams.json") if hparams.ard: # Hardcoded to WRN output size ard = 640 else: ard = None feature_extractor = WideResNet( spectral_normalization=hparams.spectral_normalization, dropout_rate=hparams.dropout_rate, coeff=hparams.coeff, n_power_iterations=hparams.n_power_iterations, batchnorm_momentum=hparams.batchnorm_momentum, ) initial_inducing_points, initial_lengthscale = initial_values_for_GP( train_dataset, feature_extractor, hparams.n_inducing_points ) gp = GP( num_outputs=num_classes, initial_lengthscale=initial_lengthscale, initial_inducing_points=initial_inducing_points, separate_inducing_points=hparams.separate_inducing_points, kernel=hparams.kernel, ard=ard, lengthscale_prior=hparams.lengthscale_prior, ) model = DKL_GP(feature_extractor, gp) model = model.cuda() likelihood = SoftmaxLikelihood(num_classes=num_classes, mixing_weights=False) likelihood = likelihood.cuda() elbo_fn = VariationalELBO(likelihood, gp, num_data=len(train_dataset)) parameters = [ {"params": feature_extractor.parameters(), "lr": hparams.learning_rate}, {"params": gp.parameters(), "lr": hparams.learning_rate}, {"params": likelihood.parameters(), "lr": hparams.learning_rate}, ] optimizer = torch.optim.SGD( parameters, momentum=0.9, weight_decay=hparams.weight_decay ) milestones = [60, 120, 160] scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=milestones, gamma=0.2 ) def step(engine, batch): model.train() likelihood.train() optimizer.zero_grad() x, y = batch x, y = x.cuda(), y.cuda() y_pred = model(x) elbo = -elbo_fn(y_pred, y) elbo.backward() optimizer.step() return elbo.item() def eval_step(engine, batch): model.eval() likelihood.eval() x, y = batch x, y = x.cuda(), y.cuda() with torch.no_grad(): y_pred = model(x) return y_pred, y trainer = Engine(step) evaluator = Engine(eval_step) metric = Average() metric.attach(trainer, "elbo") def output_transform(output): y_pred, y = output # Sample softmax values independently for classification at test time y_pred = y_pred.to_data_independent_dist() # The mean here is over likelihood samples y_pred = likelihood(y_pred).probs.mean(0) return y_pred, y metric = Accuracy(output_transform=output_transform) metric.attach(evaluator, "accuracy") metric = Loss(lambda y_pred, y: -elbo_fn(y_pred, y)) metric.attach(evaluator, "elbo") kwargs = {"num_workers": 4, "pin_memory": True} train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=hparams.batch_size, shuffle=True, drop_last=True, **kwargs, ) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=512, shuffle=False, **kwargs ) @trainer.on(Events.EPOCH_COMPLETED) def log_results(trainer): metrics = trainer.state.metrics elbo = metrics["elbo"] print(f"Train - Epoch: {trainer.state.epoch} ELBO: {elbo:.2f} ") writer.add_scalar("Likelihood/train", elbo, trainer.state.epoch) if hparams.spectral_normalization: for name, layer in model.feature_extractor.named_modules(): if isinstance(layer, torch.nn.Conv2d): writer.add_scalar( f"sigma/{name}", layer.weight_sigma, trainer.state.epoch ) if not hparams.ard: # Otherwise it's too much to submit to tensorboard length_scales = model.gp.covar_module.base_kernel.lengthscale.squeeze() for i in range(length_scales.shape[0]): writer.add_scalar( f"length_scale/{i}", length_scales[i], trainer.state.epoch ) if trainer.state.epoch > 150 and trainer.state.epoch % 5 == 0: _, auroc, aupr = get_ood_metrics( hparams.dataset, "SVHN", model, likelihood, hparams.data_root ) print(f"OoD Metrics - AUROC: {auroc}, AUPR: {aupr}") writer.add_scalar("OoD/auroc", auroc, trainer.state.epoch) writer.add_scalar("OoD/auprc", aupr, trainer.state.epoch) evaluator.run(test_loader) metrics = evaluator.state.metrics acc = metrics["accuracy"] elbo = metrics["elbo"] print( f"Test - Epoch: {trainer.state.epoch} " f"Acc: {acc:.4f} " f"ELBO: {elbo:.2f} " ) writer.add_scalar("Likelihood/test", elbo, trainer.state.epoch) writer.add_scalar("Accuracy/test", acc, trainer.state.epoch) scheduler.step() pbar = ProgressBar(dynamic_ncols=True) pbar.attach(trainer) trainer.run(train_loader, max_epochs=200) # Done training - time to evaluate results = {} evaluator.run(train_loader) train_acc = evaluator.state.metrics["accuracy"] train_elbo = evaluator.state.metrics["elbo"] results["train_accuracy"] = train_acc results["train_elbo"] = train_elbo evaluator.run(test_loader) test_acc = evaluator.state.metrics["accuracy"] test_elbo = evaluator.state.metrics["elbo"] results["test_accuracy"] = test_acc results["test_elbo"] = test_elbo _, auroc, aupr = get_ood_metrics( hparams.dataset, "SVHN", model, likelihood, hparams.data_root ) results["auroc_ood_svhn"] = auroc results["aupr_ood_svhn"] = aupr print(f"Test - Accuracy {results['test_accuracy']:.4f}") results_json = json.dumps(results, indent=4, sort_keys=True) (results_dir / "results.json").write_text(results_json) torch.save(model.state_dict(), results_dir / "model.pt") torch.save(likelihood.state_dict(), results_dir / "likelihood.pt") writer.close()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED gpus = list(config.GPUS) dataset_type = get_dataset(config) train_data = dataset_type(config, is_train=True) train_loader = DataLoader(dataset=train_data, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) val_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) # config.MODEL.NUM_JOINTS = train_data.get_num_points() model = models.get_face_alignment_net(config) # copy model files writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } model = nn.DataParallel(model, device_ids=gpus).cuda() # loss criterion = torch.nn.MSELoss(size_average=True).cuda() optimizer = utils.get_optimizer(config, model) best_nme = 100 last_epoch = config.TRAIN.BEGIN_EPOCH if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'latest.pth') if os.path.islink(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_nme = checkpoint['best_nme'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) else: print("=> no checkpoint found") for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() function.train(config, train_loader, model, criterion, optimizer, epoch, writer_dict) # evaluate nme, predictions = function.validate(config, val_loader, model, criterion, epoch, writer_dict) is_best = nme < best_nme best_nme = min(nme, best_nme) logger.info('=> saving checkpoint to {}'.format(final_output_dir)) print("best:", is_best) utils.save_checkpoint( { "state_dict": model, "epoch": epoch + 1, "best_nme": best_nme, "optimizer": optimizer.state_dict(), }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch)) final_model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) writer_dict['writer'].close()
def train(self): self.info() def KLloss(mu, logvar): return torch.mean( -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp(), dim=1), dim=0) opt = self.opt self.z = prepare_z_dist(opt.training.batch_size, opt.GenModel.style_dim, self.device, seed=self.opt.seed) self.y = prepare_y_dist(opt.training.batch_size, len(self.lexicon), self.device, seed=self.opt.seed) self.eval_z = prepare_z_dist(opt.training.eval_batch_size, opt.GenModel.style_dim, self.device, seed=self.opt.seed) self.eval_y = prepare_y_dist(opt.training.eval_batch_size, len(self.lexicon), self.device, seed=self.opt.seed) self.train_loader = DataLoader(get_dataset(opt.dataset, opt.training.dset_split), batch_size=opt.training.batch_size, shuffle=True, collate_fn=self.collect_fn, num_workers=4, drop_last=True) self.tst_loader = DataLoader( get_dataset(opt.dataset, opt.valid.dset_split), batch_size=opt.training.eval_batch_size // 2, shuffle=True, collate_fn=self.collect_fn) self.tst_loader2 = DataLoader( get_dataset(opt.dataset, opt.training.dset_split), batch_size=opt.training.eval_batch_size // 2, shuffle=True, collate_fn=self.collect_fn, num_workers=4) self.optimizers = Munch( G=torch.optim.Adam(chain(self.models.G.parameters(), self.models.E.parameters()), lr=opt.training.lr, betas=(opt.training.adam_b1, opt.training.adam_b2)), D=torch.optim.Adam(chain(self.models.D.parameters(), self.models.R.parameters(), self.models.W.parameters()), lr=opt.training.lr, betas=(opt.training.adam_b1, opt.training.adam_b2)), ) self.lr_schedulers = Munch(G=get_scheduler(self.optimizers.G, opt.training), D=get_scheduler(self.optimizers.D, opt.training)) self.averager_meters = AverageMeterManager([ 'adv_loss', 'fake_disc_loss', 'real_disc_loss', 'info_loss', 'fake_ctc_loss', 'real_ctc_loss', 'fake_wid_loss', 'real_wid_loss', 'kl_loss', 'gp_ctc', 'gp_info', 'gp_wid' ]) device = self.device ctc_len_scale = 8 best_kid = np.inf iter_count = 0 for epoch in range(1, self.opt.training.epochs): for i, (imgs, img_lens, lbs, lb_lens, wids) in enumerate(self.train_loader): ############################# # Prepare inputs & Network Forward ############################# self.set_mode('train') real_imgs, real_img_lens, real_wids = imgs.to( device), img_lens.to(device), wids.to(device) real_lbs, real_lb_lens = lbs.to(device), lb_lens.to(device) ############################# # Optimizing Recognizer & Writer Identifier & Discriminator ############################# self.optimizers.D.zero_grad() set_requires_grad([self.models.G, self.models.E], False) set_requires_grad( [self.models.R, self.models.D, self.models.W], True) ### Compute CTC loss for real samples### real_ctc = self.models.R(real_imgs) real_ctc_lens = real_img_lens // ctc_len_scale real_ctc_loss = self.ctc_loss(real_ctc, real_lbs, real_ctc_lens, real_lb_lens) self.averager_meters.update('real_ctc_loss', real_ctc_loss.item()) real_wid_logits = self.models.W(real_imgs, real_img_lens) real_wid_loss = self.classify_loss(real_wid_logits, real_wids) self.averager_meters.update('real_wid_loss', real_wid_loss.item()) with torch.no_grad(): self.y.sample_() sampled_words = idx_to_words( self.y, self.lexicon, self.opt.training.capitalize_ratio) fake_lbs, fake_lb_lens = self.label_converter.encode( sampled_words) fake_lbs, fake_lb_lens = fake_lbs.to( device).detach(), fake_lb_lens.to(device).detach() self.z.sample_() fake_imgs = self.models.G(self.z, fake_lbs, fake_lb_lens) enc_styles, _, _ = self.models.E( real_imgs, real_img_lens, self.models.W.cnn_backbone, vae_mode=True) noises = torch.randn( (real_imgs.size(0), self.opt.GenModel.style_dim - self.opt.EncModel.style_dim)).float().to(device) enc_z = torch.cat([noises, enc_styles], dim=-1) style_imgs = self.models.G(enc_z, fake_lbs, fake_lb_lens) cat_fake_imgs = torch.cat([fake_imgs, style_imgs], dim=0) cat_fake_lb_lens = fake_lb_lens.repeat(2, ).detach() cat_fake_img_lens = cat_fake_lb_lens * self.opt.char_width ### Compute discriminative loss for real & fake samples ### fake_disc = self.models.D(cat_fake_imgs.detach(), cat_fake_img_lens, cat_fake_lb_lens) fake_disc_loss = torch.mean(F.relu(1.0 + fake_disc)) real_disc = self.models.D(real_imgs, real_img_lens, real_lb_lens) real_disc_loss = torch.mean(F.relu(1.0 - real_disc)) disc_loss = real_disc_loss + fake_disc_loss self.averager_meters.update('real_disc_loss', real_disc_loss.item()) self.averager_meters.update('fake_disc_loss', fake_disc_loss.item()) (real_ctc_loss + disc_loss + real_wid_loss).backward() self.optimizers.D.step() ############################# # Optimizing Generator ############################# if iter_count % self.opt.training.num_critic_train == 0: self.optimizers.G.zero_grad() set_requires_grad( [self.models.D, self.models.R, self.models.W], False) set_requires_grad([self.models.G, self.models.E], True) ########################## # Prepare Fake Inputs ########################## self.y.sample_() sampled_words = idx_to_words( self.y, self.lexicon, self.opt.training.capitalize_ratio) fake_lbs, fake_lb_lens = self.label_converter.encode( sampled_words) fake_lbs, fake_lb_lens = fake_lbs.to( device).detach(), fake_lb_lens.to(device).detach() fake_img_lens = fake_lb_lens * self.opt.char_width self.z.sample_() fake_imgs = self.models.G(self.z, fake_lbs, fake_lb_lens) enc_styles, enc_mu, enc_logvar = self.models.E( real_imgs, real_img_lens, self.models.W.cnn_backbone, vae_mode=True) noises = torch.randn( (real_imgs.size(0), self.opt.GenModel.style_dim - self.opt.EncModel.style_dim)).float().to(device) enc_z = torch.cat([noises, enc_styles], dim=-1) style_imgs = self.models.G(enc_z, fake_lbs, fake_lb_lens) style_img_lens = fake_lb_lens * self.opt.char_width ### Concatenating all generated images in a batch ### cat_fake_imgs = torch.cat([fake_imgs, style_imgs], dim=0) cat_fake_lbs = fake_lbs.repeat(2, 1).detach() cat_fake_lb_lens = fake_lb_lens.repeat(2, ).detach() cat_fake_img_lens = cat_fake_lb_lens * self.opt.char_width ################################################### # Calculating G Losses #################################################### ### deal with fake samples ### ### Compute Adversarial loss ### cat_fake_disc = self.models.D(cat_fake_imgs, cat_fake_img_lens, cat_fake_lb_lens) adv_loss = -torch.mean(cat_fake_disc) ### CTC Auxiliary loss ### cat_fake_ctc = self.models.R(cat_fake_imgs) cat_fake_ctc_lens = cat_fake_img_lens // ctc_len_scale fake_ctc_loss = self.ctc_loss(cat_fake_ctc, cat_fake_lbs, cat_fake_ctc_lens, cat_fake_lb_lens) ### Latent Style Reconstruction ### styles = self.models.E(fake_imgs, fake_img_lens, self.models.W.cnn_backbone) info_loss = torch.mean( torch.abs( styles - self.z[:, -self.opt.EncModel.style_dim:].detach())) ### Writer Identify Loss ### recn_wid_logits = self.models.W(style_imgs, style_img_lens) fake_wid_loss = self.classify_loss(recn_wid_logits, real_wids) ### KL-Divergence Loss ### kl_loss = KLloss(enc_mu, enc_logvar) ### Gradient balance ### grad_fake_adv = torch.autograd.grad(adv_loss, cat_fake_imgs, create_graph=True, retain_graph=True)[0] grad_fake_OCR = torch.autograd.grad(fake_ctc_loss, cat_fake_ctc, create_graph=True, retain_graph=True)[0] grad_fake_info = torch.autograd.grad(info_loss, fake_imgs, create_graph=True, retain_graph=True)[0] grad_fake_wid = torch.autograd.grad(fake_wid_loss, recn_wid_logits, create_graph=True, retain_graph=True)[0] std_grad_adv = torch.std(grad_fake_adv) gp_ctc = torch.div( std_grad_adv, torch.std(grad_fake_OCR) + 1e-8).detach() + 1 gp_info = torch.div( std_grad_adv, torch.std(grad_fake_info) + 1e-8).detach() + 1 gp_wid = torch.div( std_grad_adv, torch.std(grad_fake_wid) + 1e-8).detach() + 1 self.averager_meters.update('gp_ctc', gp_ctc.item()) self.averager_meters.update('gp_info', gp_info.item()) self.averager_meters.update('gp_wid', gp_wid.item()) g_loss = 2 * adv_loss + \ gp_ctc * fake_ctc_loss + \ gp_info * info_loss + \ gp_wid * fake_wid_loss + \ self.opt.training.lambda_kl * kl_loss g_loss.backward() self.averager_meters.update('adv_loss', adv_loss.item()) self.averager_meters.update('fake_ctc_loss', fake_ctc_loss.item()) self.averager_meters.update('info_loss', info_loss.item()) self.averager_meters.update('fake_wid_loss', fake_wid_loss.item()) self.averager_meters.update('kl_loss', kl_loss.item()) self.optimizers.G.step() if iter_count % self.opt.training.print_iter_val == 0: meter_vals = self.averager_meters.eval_all() self.averager_meters.reset_all() info = "[%3d|%3d]-[%4d|%4d] G:%.4f D-fake:%.4f D-real:%.4f " \ "CTC-fake:%.4f CTC-real:%.4f Wid-fake:%.4f Wid-real:%.4f " \ "Recn-z:%.4f Kl:%.4f" \ % (epoch, self.opt.training.epochs, iter_count % len(self.train_loader), len(self.train_loader), meter_vals['adv_loss'], meter_vals['fake_disc_loss'], meter_vals['real_disc_loss'], meter_vals['fake_ctc_loss'], meter_vals['real_ctc_loss'], meter_vals['fake_wid_loss'], meter_vals['real_wid_loss'], meter_vals['info_loss'], meter_vals['kl_loss']) self.print(info) if self.writer: for key, val in meter_vals.items(): self.writer.add_scalar('loss/%s' % key, val, iter_count + 1) if (iter_count + 1) % self.opt.training.sample_iter_val == 0: if not (self.logger and self.writer): self.create_logger() sample_root = os.path.join(self.log_root, self.opt.training.sample_dir) if not os.path.exists(sample_root): os.makedirs(sample_root) self.sample_images(iter_count + 1) iter_count += 1 if epoch: ckpt_root = os.path.join(self.log_root, self.opt.training.ckpt_dir) if not os.path.exists(ckpt_root): os.makedirs(ckpt_root) self.save('last', epoch) if epoch >= self.opt.training.start_save_epoch_val and \ epoch % self.opt.training.save_epoch_val == 0: self.print('Calculate FID_KID') scores = self.validate() fid, kid = scores['FID'], scores['KID'] self.print('FID:{} KID:{}'.format(fid, kid)) if kid < best_kid: best_kid = kid self.save('best', epoch, KID=kid, FID=fid) if self.writer: self.writer.add_scalar('valid/FID', fid, epoch) self.writer.add_scalar('valid/KID', kid, epoch) for scheduler in self.lr_schedulers.values(): scheduler.step(epoch)
def main(): # args = parse_args() # logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') # logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) # gpus = list(config.GPUS) # # # load model state_dict = torch.load(args.model_file) if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict'] model.load_state_dict(state_dict) else: model.module.load_state_dict(state_dict) model = nn.DataParallel(model, device_ids=gpus).cuda() # dataset_type = get_dataset(config) test_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) nme, predictions = function.inference(config, test_loader, model) torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth')) target = test_loader.dataset.load_all_pts() pred = 16 * predictions l = len(pred) res = 0.0 res_tmp = [0.0 for i in range(config.MODEL.NUM_JOINTS)] res_tmp = np.array(res_tmp) res_temp_x = target - pred res_temp_x = res_temp_x[:, :, 0] res_temp_y = target - pred res_temp_y = res_temp_y[:, :, 1] # csv_file_test_x = pd.DataFrame(np.transpose(np.array(pred[:, :, 0])), columns=test_loader.dataset.annotation_files) # csv_file_test_y = pd.DataFrame(np.transpose(np.array(pred[:, :, 1])), columns=test_loader.dataset.annotation_files) # csv_file_target_x = pd.DataFrame(np.transpose(np.array(target[:, :, 0])), columns=test_loader.dataset.annotation_files) # csv_file_target_y = pd.DataFrame(np.transpose(np.array(target[:, :, 1])), columns=test_loader.dataset.annotation_files) for i in range(l): trans = np.sqrt( pow(target[i][0][0] - target[i][1][0], 2) + pow(target[i][0][1] - target[i][1][1], 2)) / 30.0 res_temp_x[i] = res_temp_x[i] / trans res_temp_y[i] = res_temp_y[i] / trans for j in range(len(target[i])): dist = np.sqrt( np.power((target[i][j][0] - pred[i][j][0]), 2) + np.power((target[i][j][1] - pred[i][j][1]), 2)) / trans res += dist res_tmp[j] += dist res_t = np.sqrt(res_temp_x * res_temp_x + res_temp_y * res_temp_y) # pd.DataFrame(data=res_temp_x.data.value).to_csv('res_x') # pd.DataFrame(data=res_temp_y.data.value).to_csv('res_y') # pd.DataFrame(data=res_t.data.value).to_csv('res_t') res_tmp /= np.float(len(pred)) print(res_tmp) print(np.mean(res_tmp)) res /= (len(pred) * len(pred[0])) print(res)
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn.benchmark = config.CUDNN.BENCHMARK # cudnn.determinstic = config.CUDNN.DETERMINISTIC # cudnn.enabled = config.CUDNN.ENABLED # if isinstance(config.TRAIN.LR_STEP, list): # lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( # optimizer, config.TRAIN.LR_STEP, # # config.TRAIN.LR_FACTOR, last_epoch-1 # config.TRAIN.LR_FACTOR, 0 # ) # else: # lr_scheduler = torch.optim.lr_scheduler.StepLR( # optimizer, config.TRAIN.LR_STEP, # # config.TRAIN.LR_FACTOR, last_epoch-1 # config.TRAIN.LR_FACTOR, 0 # ) dataset_type = get_dataset(config) train_dataset = dataset_type(config, is_train=True) # train_dataset[0] # return 0 train_loader = DataLoader( dataset=dataset_type(config, is_train=True), # batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus), batch_size=config.TRAIN.BATCH_SIZE_PER_GPU, shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS) # val_loader = DataLoader( # dataset=dataset_type(config, # is_train=True), # # batch_size=config.TEST.BATCH_SIZE_PER_GPU*len(gpus), # batch_size=config.TEST.BATCH_SIZE_PER_GPU, # shuffle=False, # num_workers=config.WORKERS, # # pin_memory=config.PIN_MEMORY # ) model = models.get_face_alignment_net(config) # copy model files writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } gpus = list(config.GPUS) # model = nn.DataParallel(model, device_ids=gpus).cuda() model.to("cuda") # loss criterion = torch.nn.MSELoss(size_average=True).cuda() # criterion = fnn.mse_loss # criterion = WingLoss() # criterion = Loss_weighted() optimizer = utils.get_optimizer(config, model) best_nme = 100 last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'latest.pth') if os.path.isfile(model_state_file): with open(model_state_file, "rb") as fp: state_dict = torch.load(fp) model.load_state_dict(state_dict) last_epoch = 1 # checkpoint = torch.load(model_state_file) # last_epoch = checkpoint['epoch'] # best_nme = checkpoint['best_nme'] # model.load_state_dict(checkpoint['state_dict']) # optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint (epoch {})" .format(last_epoch)) else: print("=> no checkpoint found") lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() function.train(config, train_loader, model, criterion, optimizer, epoch, writer_dict) # evaluate nme = 0 # nme, predictions = function.validate(config, val_loader, model, # criterion, epoch, writer_dict) is_best = True # is_best = nme < best_nme best_nme = min(nme, best_nme) logger.info('=> saving checkpoint to {}'.format(final_output_dir)) print("best:", is_best) torch.save(model.state_dict(), os.path.join(final_output_dir, 'mse_relu_checkpoint_{}.pth'.format(epoch))) # utils.save_checkpoint( # {"state_dict": model, # "epoch": epoch + 1, # "best_nme": best_nme, # "optimizer": optimizer.state_dict(), # }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch)) final_model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info('saving final model state to {}'.format( final_model_state_file)) torch.save(model.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): args = parse_args() # set logger and dir logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.experiment_name, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # set cudnn cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # 目前仅支持单gpu,todo:增加多gpu支持 # set model and loss and criterion model = models.get_face_alignment_net(config) model = model.cuda(config.GPUS[0]) criterion = torch.nn.MSELoss(size_average=True).cuda(config.GPUS[0]) # criterion = AdaptiveWingLoss() optimizer = utils.get_optimizer(config, model) # get dataset dataset_type = get_dataset(config) # get dataloader train_loader = DataLoader(dataset=dataset_type(config, is_train=True), batch_size=config.TRAIN.BATCH_SIZE_PER_GPU, shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) val_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU, shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) # set lr_scheduler last_epoch = config.TRAIN.BEGIN_EPOCH if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) # set training writer writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } # set training resume function if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'latest.pth') if os.path.islink(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_nme = checkpoint['best_nme'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) else: print("=> no checkpoint found") # starting training best_nme = 10000 for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() # traing function.train(config, train_loader, model, criterion, optimizer, epoch, writer_dict) # evaluating nme, predictions = function.validate(config, val_loader, model, criterion, epoch, writer_dict) # saving is_best = nme < best_nme best_nme = min(nme, best_nme) logger.info('=> saving checkpoint to {}'.format(final_output_dir)) print("best:", is_best) utils.save_checkpoint( { "state_dict": model, "epoch": epoch + 1, "best_nme": best_nme, "optimizer": optimizer.state_dict(), }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch)) final_model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main_wflw(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) dataset_type = get_dataset(config) dataset = dataset_type(config, is_train=True) fp = open('data/wflw/face_landmarks70_wflw_train.csv', 'w') for i in range(len(dataset)): img, image_path, meta = dataset[i] fold, name = image_path.split('/')[-2], image_path.split('/')[-1] folder = osp.join('data/wflw/xximages', fold) if not osp.exists(folder): os.makedirs(folder) fname = osp.join(folder, name) scale = meta['scale'] center = meta['center'] tpts = meta['tpts'] selpts = [] for j in range(0, 33, 2): selpts.append(tpts[j]) # eyebow selpts.append(tpts[33]) selpts.append((tpts[34] + tpts[41]) / 2) selpts.append((tpts[35] + tpts[40]) / 2) selpts.append((tpts[36] + tpts[39]) / 2) selpts.append((tpts[37] + tpts[38]) / 2) selpts.append((tpts[42] + tpts[50]) / 2) selpts.append((tpts[43] + tpts[49]) / 2) selpts.append((tpts[44] + tpts[48]) / 2) selpts.append((tpts[45] + tpts[47]) / 2) selpts.append(tpts[46]) # nose for j in range(51, 60): selpts.append(tpts[j]) # eye selpts.append(tpts[60]) selpts.append((tpts[61] + tpts[62]) / 2) selpts.append(tpts[63]) selpts.append(tpts[64]) selpts.append(tpts[65]) selpts.append((tpts[66] + tpts[67]) / 2) selpts.append(tpts[68]) selpts.append(tpts[69]) selpts.append((tpts[70] + tpts[71]) / 2) selpts.append(tpts[72]) selpts.append((tpts[73] + tpts[74]) / 2) selpts.append(tpts[75]) for j in range(76, 98): selpts.append(tpts[j]) fp.write('%s,%.2f,%.1f,%.1f' % (osp.join(fold, name), scale, center[0], center[1])) for spt in selpts: cv2.circle(img, (spt[0], spt[1]), 1, (0, 0, 255)) fp.write(',%f,%f' % (spt[0], spt[1])) fp.write('\n') img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite(fname, img) fp.close()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn.benchmark = config.CUDNN.BENCHMARK # cudnn.determinstic = config.CUDNN.DETERMINISTIC # cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) gpus = list(config.GPUS) # model = nn.DataParallel(model, device_ids=gpus).cuda() model.to("cuda") # print(model) # load model # state_dict = torch.load(args.model_file) # print(state_dict) # model = torch.load(args.model_file) with open(args.model_file, "rb") as fp: state_dict = torch.load(fp) model.load_state_dict(state_dict) # model.load_state_dict(state_dict['state_dict']) # if 'state_dict' in state_dict.keys(): # state_dict = state_dict['state_dict'] # # print(state_dict) # model.load_state_dict(state_dict) # else: # model.module.load_state_dict(state_dict) dataset_type = get_dataset(config) test_loader = DataLoader( dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU*len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY ) predictions = function.inference(config, test_loader, model) # print("len(predictions)", len(predictions)) # print(predictions[0]) df_predictions = [] for pred in predictions: row = dict() row['file_name'] = pred[0] for id_point in range(194): row[f'Point_M{id_point}_X'] = int(pred[1][id_point]) row[f'Point_M{id_point}_Y'] = int(pred[2][id_point]) df_predictions.append(row) df_predictions = pd.DataFrame(df_predictions) # print(predictions_meta[0]) df_predictions.to_csv('pred_test.csv', index=False)
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED gpus = list(config.GPUS) dataset_type = get_dataset(config) train_data = dataset_type(config, split="train") train_loader = DataLoader(dataset=train_data, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) val_data = dataset_type(config, split="valid") val_loader = DataLoader(dataset=val_data, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) # config.MODEL.NUM_JOINTS = train_data.get_num_points() model = models.get_face_alignment_net(config) # copy model files writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } model = nn.DataParallel(model, device_ids=gpus).cuda() # loss criterion = torch.nn.MSELoss(size_average=True).cuda() optimizer = utils.get_optimizer(config, model) best_nme = 100 last_epoch = config.TRAIN.BEGIN_EPOCH if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'final.pth') if os.path.islink(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_nme = checkpoint['best_nme'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) else: print("=> no checkpoint found") loss = [] for epoch in range(last_epoch, config.TRAIN.END_EPOCH): losses, diff = function.train(config, train_loader, model, criterion, optimizer, epoch, writer_dict) loss.append(losses) lr_scheduler.step() np.save( os.path.join(final_output_dir, "train_diff@epoch{}".format(epoch)), diff) # evaluate nme, predictions, diff = function.validate(config, val_loader, model, criterion, epoch, writer_dict) np.save( os.path.join(final_output_dir, "valid_diff@epoch{}".format(epoch)), diff) is_best = nme < best_nme best_nme = min(nme, best_nme) logger.info('=> saving checkpoint to {}'.format(final_output_dir)) print("best:", is_best) utils.save_checkpoint( { "state_dict": model, "epoch": epoch + 1, "best_nme": best_nme, "optimizer": optimizer.state_dict(), }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch)) if is_best: for i in range(len(predictions)): afile = val_data.annotation_files[i] new_afile = '{}.{}.txt'.format( afile, os.path.basename(args.cfg).split('.')[0]) with open(new_afile, 'wt') as f: pts = predictions[i].cpu().numpy() for j in range(len(pts)): f.write("{},{}\n".format( pts[j][1] / val_data.factor[1], pts[j][0] / val_data.factor[0])) pd.DataFrame(data=loss).to_csv('loss2.csv') final_model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) writer_dict['writer'].close()