def srpgan_eval(images, generator_file, upscale_factor, use_cuda, num_workers=os.cpu_count()): """ Turn a list of images to super resolution and returns them Args: num_workers (int): Number of processors to use use_cuda (bool): Whether or not to use the GPU upscale_factor (int): Either 2, 4 or 8 images (list): List of Pillow images generator_file (file): The generator saved model file Returns: list: A list of SR images """ netG = Generator(upscale_factor) learner = Learner(ClassifierCore(netG, None, None), use_cuda=use_cuda) ModelSaverCallback.restore_model_from_file(netG, generator_file, load_with_cpu=not use_cuda) eval_ds = EvalDataset(images) # One batch at a time as the pictures may differ in size eval_dl = DataLoader(eval_ds, 1, shuffle=False, num_workers=num_workers) images_pred = [] predictions = learner.predict(eval_dl, flatten_predictions=False) tfs = transforms.Compose([ transforms.ToPILImage(), ]) for pred in predictions: pred = pred.view(pred.size()[1:]) # Remove batch size == 1 images_pred.append(tfs(pred.cpu())) return images_pred
def train(args): num_workers = os.cpu_count() train_loader, valid_loader = get_loaders(args, num_workers) model_saver = ModelSaverCallback(saved_model_dir.absolute(), args.adv_epochs, every_n_epoch=10) netG = Generator(args.upscale_factor) netG.apply(weights_init) netD = Discriminator((3, args.crop_size, args.crop_size)) netD.apply(weights_init) optimizer_g = optim.Adam(netG.parameters(), lr=1e-4) optimizer_d = optim.Adam(netD.parameters(), lr=1e-4) # Restore models if they exists if args.restore_models == 1: model_saver.restore_models([netG, netD], saved_model_dir.absolute()) else: if args.gen_epochs > 0: print("---------------------- Generator training ----------------------") callbacks = [ReduceLROnPlateau(optimizer_g, loss_step="train")] loss = nn.MSELoss() learner = Learner(ClassifierCore(netG, optimizer_g, loss)) learner.train(args.gen_epochs, None, train_loader, None, callbacks) print("----------------- Adversarial (SRPGAN) training -----------------") callbacks = [model_saver, ReduceLROnPlateau(optimizer_g, loss_step="valid"), TensorboardVisualizerCallback(tensorboard_dir.absolute())] g_loss = GeneratorLoss() d_loss = DiscriminatorLoss() learner = Learner(SRPGanCore(netG, netD, optimizer_g, optimizer_d, g_loss, d_loss)) learner.train(args.adv_epochs, [SSIM(), PSNR()], train_loader, valid_loader, callbacks)
vectorizer=vectorizer, ) #loss = torch.nn.MSELoss() loss = torch.nn.CosineEmbeddingLoss(margin=0.2) model = SiamesCos( word_emb_sizes=args.emb_size, conv_sizes=[int(args.netsize)], out_size=[args.netsize, args.netsize // 2], dropout=args.dropout, ) if args.restore_model: ModelSaverCallback.restore_model_from_file(model, args.restore_model, load_with_cpu=(not args.cuda)) optimizer = optim.Adam(filter(lambda x: x.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) run_name = args.run + '-' + datetime.datetime.now().strftime("%Y-%m-%d-%H-%M") tb_dir = os.path.join(TB_DIR, run_name) if not os.path.exists(tb_dir): os.mkdir(tb_dir) metrics = [ CosRocAucMetric(threshold=0) ]
if args.emb_path: vectorizer = ModelVectorizer(model_path=args.emb_path) else: vectorizer = OnDiskVectorizer(mtx_path=args.mtx_path, meta_path=args.meta_path) model = Matcher( word_emb_sizes=net_params['word_emb_sizes'], conv_sizes=net_params['conv_sizes'], out_size=net_params['out_size'], dropout=0.0, ) ModelSaverCallback.restore_model_from_file(model, args.model, load_with_cpu=(not args.cuda)) model.eval() with open(args.input_file) as fd: for line in fd: line = line.strip() sent_a = line sent_a_token = tokenizer(sent_a) if len(sent_a_token) <= 2: sent_a_token += ['xxx'] sent_a_vect = vectorizer.convert([sent_a_token, sent_a_token]) vect = model.cdssm_a.process_sentences(sent_a_vect)
from torchlite.torch.models.srpgan import Generator from torchlite.torch.learner import Learner from torchlite.torch.learner.cores import ClassifierCore from torchlite.torch.train_callbacks import ModelSaverCallback import os import torchvision.transforms as transforms from torch.utils.data import DataLoader cur_path = os.path.dirname(os.path.abspath(__file__)) saved_model_dir = efiles.create_dir_if_not_exists( os.path.join(cur_path, "models")) netG = Generator(4) learner = Learner(ClassifierCore(netG, None, None), use_cuda=True) ModelSaverCallback.restore_model_from_file(netG, saved_model_dir / "Generator.pth", load_with_cpu=not True) def evaluate(imorig): """ Method used for inference only """ img_shape = imorig.shape h = 200 num_h = img_shape[0] // h ori_h = h * num_h crop_h = img_shape[0] - ori_h w = 200 num_w = img_shape[1] // w ori_w = w * num_w