def test_random_choice(self): random_state = random.getstate() random.seed(42) random_choice_transform = transform.RandomChoice([ transform.Resize(15), transform.Resize(20), transform.CenterCrop(10) ]) img = transform.ToPILImage()(jt.random((25, 25, 3))) num_samples = 250 num_resize_15 = 0 num_resize_20 = 0 num_crop_10 = 0 for _ in range(num_samples): out = random_choice_transform(img) if out.size == (15, 15): num_resize_15 += 1 elif out.size == (20, 20): num_resize_20 += 1 elif out.size == (10, 10): num_crop_10 += 1 p_value = stats.binom_test(num_resize_15, num_samples, p=0.33333) self.assertGreater(p_value, 0.0001) p_value = stats.binom_test(num_resize_20, num_samples, p=0.33333) self.assertGreater(p_value, 0.0001) p_value = stats.binom_test(num_crop_10, num_samples, p=0.33333) self.assertGreater(p_value, 0.0001) random.setstate(random_state)
def __init__(self, root, lmdir, maskdir, cmaskdir, mode="test", load_h=512, load_w=512): super().__init__() transform_ = [ transform.Resize((load_h, load_w), Image.BICUBIC), transform.ImageNormalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ] self.transform = transform.Compose(transform_) transform_mask_ = [ transform.Resize((load_h, load_w), Image.BICUBIC), transform.Gray(), ] self.transform_mask = transform.Compose(transform_mask_) self.files_A = sorted(glob.glob(root + "/*.*")) self.total_len = len(self.files_A) self.batch_size = None self.shuffle = False self.drop_last = False self.num_workers = None self.buffer_size = 512*1024*1024 self.lmdir = lmdir self.maskdir = maskdir self.cmaskdir = cmaskdir self.load_h = load_h
def main(): parser = argparse.ArgumentParser() parser.add_argument('--batch_size', type=int, default=32) parser.add_argument('--epochs', type=int, default=50) parser.add_argument('--num_classes', type=int, default=130) parser.add_argument('--lr', type=float, default=2e-3) parser.add_argument('--weight_decay', type=float, default=1e-5) parser.add_argument('--resume', type=bool, default=False) parser.add_argument('--eval', type=bool, default=False) parser.add_argument('--dataroot', type=str, default='/content/drive/MyDrive/dogflg/data2/') parser.add_argument('--model_path', type=str, default='./best_model.bin') parser.add_argument('--sampleratio', type=float, default=0.8) args = parser.parse_args() transform_train = transform.Compose([ transform.Resize((256, 256)), transform.CenterCrop(224), transform.RandomHorizontalFlip(), transform.ToTensor(), transform.ImageNormalize(0.485, 0.229), # transform.ImageNormalize(0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) root_dir = args.dataroot train_loader = TsinghuaDog(root_dir, batch_size=args.batch_size, train=True, part='train', shuffle=True, transform=transform_train, sample_rate=args.sampleratio) transform_test = transform.Compose([ transform.Resize((256, 256)), transform.CenterCrop(224), transform.ToTensor(), transform.ImageNormalize(0.485, 0.229), # transform.ImageNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) val_loader = TsinghuaDog(root_dir, batch_size=args.batch_size, train=False, part='val', shuffle=False, transform=transform_test, sample_rate=args.sampleratio) epochs = args.epochs model = Net(num_classes=args.num_classes) lr = args.lr weight_decay = args.weight_decay optimizer = SGD(model.parameters(), lr=lr, momentum=0.99) if args.resume: model.load(args.model_path) print('model loaded', args.model_path) #random save for test #model.save(args.model_path) if args.eval: evaluate(model, val_loader, save_path=args.model_path) return for epoch in range(epochs): train(model, train_loader, optimizer, epoch) evaluate(model, val_loader, epoch, save_path=args.model_path)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--batch_size', type=int, default=8) parser.add_argument('--epochs', type=int, default=50) parser.add_argument('--num_classes', type=int, default=130) parser.add_argument('--lr', type=float, default=2e-3) parser.add_argument('--weight_decay', type=float, default=1e-5) parser.add_argument('--resume', type=bool, default=False) parser.add_argument('--eval', type=bool, default=False) parser.add_argument('--dataroot', type=str, default='/home/gmh/dataset/TsinghuaDog/') parser.add_argument('--model_path', type=str, default='./best_model.pkl') args = parser.parse_args() transform_train = transform.Compose([ transform.Resize((512, 512)), transform.RandomCrop(448), transform.RandomHorizontalFlip(), transform.ToTensor(), transform.ImageNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) root_dir = args.dataroot train_loader = TsinghuaDog(root_dir, batch_size=16, train=True, part='train', shuffle=True, transform=transform_train) transform_test = transform.Compose([ transform.Resize((512, 512)), transform.CenterCrop(448), transform.ToTensor(), transform.ImageNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) val_loader = TsinghuaDog(root_dir, batch_size=16, train=False, part='val', shuffle=False, transform=transform_test) epochs = args.epochs model = Net(num_classes=args.num_classes) lr = args.lr weight_decay = args.weight_decay optimizer = SGD(model.parameters(), lr=lr, momentum=0.9) if args.resume: model.load(args.model_path) if args.eval: evaluate(model, val_loader) return for epoch in range(epochs): train(model, train_loader, optimizer, epoch) evaluate(model, val_loader, epoch)
def main(): best_acc1 = 0 #TODO multi gpu if args.model == 'TNet26': model = TNet.Resnet26() elif args.model == 'TNet38': model = TNet.Resnet38() elif args.model == 'TNet50': model = TNet.Resnet50() elif args.model == 'Resnet26': model = jtmodels.__dict__['resnet26']() elif args.model == 'Resnet38': model = jtmodels.__dict__['resnet38']() elif args.model == 'Resnet50': model = jtmodels.__dict__['resnet50']() elif args.model == 'SAN10': model = san(sa_type = 0, layers=[2, 1, 2, 4, 1], kernels=[3, 7, 7, 7, 7], num_classes = 1000) elif args.model == 'SAN_TCN10': model = san_tcn(sa_type = 0, layers=[2, 1, 2, 4, 1], kernels=[3, 7, 7, 7, 7], num_classes = 1000) else: print("Model not found!") exit(0) if (args.use_pytorch_conv_init): pytorch_conv_init(model) model_path = os.path.join(args.save_path, 'model_best.pk') model.load(model_path) mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] # val_transform = transform.Compose([transform.Resize(256), transform.CenterCrop(224), transform.ImageNormalize(mean, std)]) # [transform diff from val!vvvv] val_transform = transform.Compose([transform.Resize(224), transform.ImageNormalize(mean, std)]) val_loader = ImageFolder('input_images', val_transform).set_attrs(batch_size=args.batch_size_val, shuffle=False, num_workers=args.workers) test(val_loader, model)
def get_loader(root_dir, label_file, batch_size, img_size=0, num_thread=4, pin=True, test=False, split='train'): if test is False: raise NotImplementedError else: transform = transforms.Compose([ transforms.Resize((400, 400)), transforms.ToTensor(), transforms.ImageNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) dataset = SemanLineDatasetTest(root_dir, label_file, transform=transform, t_transform=None) if test is False: raise NotImplementedError else: dataset.set_attrs(batch_size=batch_size, shuffle=False) print('Get dataset success.') return dataset
def get_transform(params, gray=False, mask=False): transform_ = [] # resize transform_.append( transform.Resize((params['load_h'], params['load_w']), Image.BICUBIC)) # flip if params['flip']: transform_.append(transform.Lambda(lambda img: transform.hflip(img))) if gray: transform_.append(transform.Gray()) if mask: transform_.append(transform.ImageNormalize([ 0., ], [ 1., ])) else: if not gray: transform_.append( transform.ImageNormalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])) else: transform_.append(transform.ImageNormalize([ 0.5, ], [ 0.5, ])) return transform.Compose(transform_)
def test_random_order(self): random_state = random.getstate() random.seed(42) random_order_transform = transform.RandomOrder( [transform.Resize(20), transform.CenterCrop(10)]) img = transform.ToPILImage()(jt.random((3, 25, 25))) num_samples = 250 num_normal_order = 0 resize_crop_out = transform.CenterCrop(10)(transform.Resize(20)(img)) for _ in range(num_samples): out = random_order_transform(img) if out == resize_crop_out: num_normal_order += 1 p_value = stats.binom_test(num_normal_order, num_samples, p=0.5) random.setstate(random_state) self.assertGreater(p_value, 0.0001)
def setUpClass(self): # hyper-parameters self.batch_size = 32 self.weight_decay = 0.0001 self.momentum = 0.9 self.learning_rate = 0.01 # mnist dataset self.train_loader = MNIST(train=True, transform=trans.Resize(224)) \ .set_attrs(batch_size=self.batch_size, shuffle=True)
def __init__(self, root, hr_shape): hr_height, hr_width = hr_shape # transform for low resolution images and high resolution images self.lr_transform = transform.Compose([ transform.Resize((hr_height // 4, hr_height // 4), Image.BICUBIC), transform.ImageNormalize(mean, std), ]) self.hr_transform = transform.Compose([ transform.Resize((hr_height, hr_height), Image.BICUBIC), transform.ImageNormalize(mean, std), ]) self.files = sorted(glob.glob(root + "/*.*")) self.total_len = len(self.files) self.batch_size = None self.shuffle = False self.drop_last = False self.num_workers = None self.buffer_size = 512 * 1024 * 1024
def get_dataset(): dataset = ImageFolder(traindir).set_attrs(batch_size=256, shuffle=False) dataset.set_attrs(transform=transform.Compose([ transform.Resize(224), transform.ImageNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]), num_workers=0) return dataset
def test_resize(self): height = random.randint(24, 32) * 2 width = random.randint(24, 32) * 2 osize = random.randint(5, 12) * 2 img = jt.ones([height, width, 3]) result = transform.Compose([ transform.ToPILImage(), transform.Resize(osize), transform.ToTensor(), ])(img) self.assertIn(osize, result.shape) if height < width: self.assertLessEqual(result.shape[1], result.shape[2]) elif width < height: self.assertGreaterEqual(result.shape[1], result.shape[2]) result = transform.Compose([ transform.ToPILImage(), transform.Resize([osize, osize]), transform.ToTensor(), ])(img) self.assertIn(osize, result.shape) self.assertEqual(result.shape[1], osize) self.assertEqual(result.shape[2], osize) oheight = random.randint(5, 12) * 2 owidth = random.randint(5, 12) * 2 result = transform.Compose([ transform.ToPILImage(), transform.Resize((oheight, owidth)), transform.ToTensor(), ])(img) self.assertEqual(result.shape[1], oheight) self.assertEqual(result.shape[2], owidth) result = transform.Compose([ transform.ToPILImage(), transform.Resize([oheight, owidth]), transform.ToTensor(), ])(img) self.assertEqual(result.shape[1], oheight) self.assertEqual(result.shape[2], owidth)
def setUpClass(self): # hyper-parameters self.batch_size = int(os.environ.get("TEST_BATCH_SIZE", "100")) self.weight_decay = 0.0001 self.momentum = 0.9 self.learning_rate = 0.1 # mnist dataset self.train_loader = MNIST(train=True, transform=trans.Resize(224)) \ .set_attrs(batch_size=self.batch_size, shuffle=True) self.train_loader.num_workers = 4
def main(): batch_size = 64 learning_rate = 0.1 momentum = 0.9 weight_decay = 1e-4 epochs = 5 train_loader = MNIST(train=True, transform=trans.Resize(28)).set_attrs( batch_size=batch_size, shuffle=True) val_loader = MNIST(train=True, transform=trans.Resize(28)).set_attrs(batch_size=1, shuffle=False) model = Model() optimizer = nn.SGD(model.parameters(), learning_rate, momentum, weight_decay) for epoch in range(epochs): train(model, train_loader, optimizer, epoch) test(model, val_loader, epoch)
def __init__(self, root, input_shape, mode="train"): self.transform = transform.Compose([ transform.Resize(input_shape[-2:]), transform.ImageNormalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) self.files = sorted(glob.glob(os.path.join(root, mode) + "/*.*")) self.total_len = len(self.files) self.batch_size = None self.shuffle = False self.drop_last = False self.num_workers = None self.buffer_size = 512 * 1024 * 1024
def main(): parser = argparse.ArgumentParser() parser.add_argument('--batch_size', type=int, default=8) parser.add_argument('--epochs', type=int, default=50) parser.add_argument('--num_classes', type=int, default=130) parser.add_argument('--lr', type=float, default=2e-3) parser.add_argument('--weight_decay', type=float, default=1e-5) parser.add_argument('--resume', type=bool, default=True) parser.add_argument('--eval', type=bool, default=False) parser.add_argument('--dataroot', type=str, default='/content/drive/MyDrive/dogfl/data/TEST_A/') parser.add_argument('--model_path', type=str, default='./best_model.bin') parser.add_argument('--out_file', type=str, default='./result.json') args = parser.parse_args() root_dir = args.dataroot transform_test = transform.Compose([ transform.Resize((512, 512)), transform.CenterCrop(448), transform.ToTensor(), transform.ImageNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) name_list = [] for _, _, _name_list in os.walk(root_dir): name_list = _name_list val_loader = TsinghuaDogExam(root_dir, batch_size=args.batch_size, train=False, name_list=name_list, shuffle=False, transform=transform_test) model = Net(num_classes=args.num_classes) if args.resume: model.load(args.model_path) print('model loaded', args.model_path) top5_class_list = evaluate(model, val_loader) # label start from 1, however it doesn't pred_result = dict(zip(name_list, top5_class_list)) with open(args.out_file, 'w') as fout: json.dump(pred_result, fout, ensure_ascii=False, indent=4)
def test_not_pil_image(self): img = jt.random((30, 40, 3)) result = transform.Compose([ transform.RandomAffine(20), transform.ToTensor(), ])(img) img = jt.random((30, 40, 3)) result = transform.Compose([ transform.ToPILImage(), transform.Gray(), transform.Resize(20), transform.ToTensor(), ])(img)
def test_dataset(self): return self.train_loader = MNIST(train=True, transform=trans.Resize(224)) \ .set_attrs(batch_size=300, shuffle=True) self.train_loader.num_workers = 1 import time for batch_idx, (data, target) in tqdm(enumerate(self.train_loader)): # time.sleep(5) # print("break") # break # self.train_loader.display_worker_status() if batch_idx > 30: break pass for batch_idx, (data, target) in tqdm(enumerate(self.train_loader)): # time.sleep(5) # print("break") # break # self.train_loader.display_worker_status() if batch_idx > 300: break pass
def get_transform(new_size=None): """ obtain the image transforms required for the input data :param new_size: size of the resized images :return: image_transform => transform object from TorchVision """ # from torchvision.transforms import ToTensor, Normalize, Compose, Resize, RandomHorizontalFlip if new_size is not None: image_transform = transform.Compose([ transform.RandomHorizontalFlip(), transform.Resize(new_size), transform.ToTensor(), transform.ImageNormalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) else: image_transform = transform.Compose([ transform.RandomHorizontalFlip(), transform.ToTensor(), transform.ImageNormalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) return image_transform
def test_dataset(): root = '/home/gmh/dataset/TsinghuaDog' part = 'train' # from torchvision import transforms rgb_mean = [0.5, 0.5, 0.5] rgb_std = [0.5, 0.5, 0.5] transform_val = transform.Compose([ transform.Resize((299, 299)), transform.ToTensor(), transform.ImageNormalize(rgb_mean, rgb_std), ]) dataloader = TsinghuaDog(root, batch_size=16, train=False, part=part, shuffle=True, transform=transform_val) # def __init__(self, root_dir, batch_size, part='train', train=True, shuffle=False, transform=None, num_workers=1): for images, labels in dataloader: # print(images.size(),labels.size(),labels) pass
return self.model(d_in) # 损失函数:平方误差 # 调用方法:adversarial_loss(网络输出A, 分类标签B) # 计算结果:(A-B)^2 adversarial_loss = nn.MSELoss() generator = Generator() discriminator = Discriminator() # 导入MNIST数据集 from jittor.dataset.mnist import MNIST import jittor.transform as transform transform = transform.Compose([ transform.Resize(opt.img_size), transform.Gray(), transform.ImageNormalize(mean=[0.5], std=[0.5]), ]) dataloader = MNIST(train=True, transform=transform).set_attrs(batch_size=opt.batch_size, shuffle=True) optimizer_G = nn.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) optimizer_D = nn.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) from PIL import Image
def main(): parser = argparse.ArgumentParser() parser.add_argument('--batch_size', type=int, default=8) parser.add_argument('--epochs', type=int, default=50) parser.add_argument('--num_classes', type=int, default=130) parser.add_argument('--lr', type=float, default=2e-3) parser.add_argument('--weight_decay', type=float, default=1e-5) parser.add_argument('--resume', type=bool, default=True) parser.add_argument('--eval', type=bool, default=False) parser.add_argument('--dataroot', type=str, default='/content/drive/MyDrive/dogfl/data/TEST_A/') parser.add_argument('--model_path', type=str, default='model/res152_8/model.bin:model/res152_10/model.bin:model/seres152_8/model.bin:model/seres152_10/model.bin') parser.add_argument('--model_name', type=str, default='Net1_z:Net1_z:Net10_z:Net10_z') parser.add_argument('--out_file', type=str, default='./result.json') args = parser.parse_args() root_dir = args.dataroot transform_test = transform.Compose([ transform.Resize((256, 256)), transform.CenterCrop(224), transform.ToTensor(), transform.ImageNormalize(0.485, 0.229), ]) model_zoo = args.model_path.split(':') model_name_class = args.model_name.split(':') print(model_zoo, model_name_class) zoo_pred_result = [] i = 0 for model_path_ in model_zoo: name_list = [] for _, _, _name_list in os.walk(root_dir): name_list = _name_list val_loader = TsinghuaDogExam(root_dir, batch_size=args.batch_size, train=False, name_list=name_list, shuffle=False, transform=transform_test) model = eval(model_name_class[i])(num_classes=args.num_classes) if args.resume: model.load(model_path_) top5_class_list = evaluate(model, val_loader) # label start from 1, however it doesn't pred_result = dict(zip(name_list, top5_class_list)) zoo_pred_result.append(pred_result) i += 1 # vote the best president = zoo_pred_result[0] vote_result = {} for key in president.keys(): val_list = [] for i in range(5): candiates = [model[key][i] for model in zoo_pred_result] val, n = Counter(candiates).most_common()[0] val_list.append(val) vote_result[key] = val_list with open(args.out_file, 'w') as fout: json.dump(vote_result, fout, ensure_ascii=False, indent=4)
out = self.model(img) out = out.view((out.shape[0], (-1))) validity = self.adv_layer(out) return validity adversarial_loss = nn.BCELoss() lambda_gp = 10 # Initialize generator and discriminator generator = Generator() discriminator = Discriminator() # Configure data loader transform = transform.Compose([ transform.Resize(size=opt.img_size), transform.Gray(), transform.ImageNormalize(mean=[0.5], std=[0.5]), ]) dataloader = MNIST(train=True, transform=transform).set_attrs(batch_size=opt.batch_size, shuffle=True) # Optimizers optimizer_G = jt.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) optimizer_D = jt.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
generator = GeneratorResNet(img_shape=img_shape, res_blocks=opt.residual_blocks, c_dim=c_dim) discriminator = Discriminator(img_shape=img_shape, c_dim=c_dim) # Optimizers optimizer_G = jt.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) optimizer_D = jt.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) # Configure dataloaders train_transforms = [ transform.Resize(int(1.12 * opt.img_height), Image.BICUBIC), transform.RandomCrop(opt.img_height), transform.RandomHorizontalFlip(), transform.ImageNormalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ] dataloader = CelebADataset("../../data/%s" % opt.dataset_name, transform_=train_transforms, mode="train", attributes=opt.selected_attrs).set_attrs( batch_size=opt.batch_size, shuffle=True, num_workers=opt.n_cpu, ) val_transforms = [ transform.Resize((opt.img_height, opt.img_width), Image.BICUBIC), transform.ImageNormalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
labels = [path.split('/')[-1].split('.')[0] for path in train_list] ## Split train_list, valid_list = train_test_split(train_list, test_size=0.2, stratify=labels, random_state=42) print(f"Train Data: {len(train_list)}") print(f"Validation Data: {len(valid_list)}") print(f"Test Data: {len(test_list)}") ## Image Augumentation transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomCropAndResize(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ]) ## Load Datasets class CatsDogsDataset(Dataset): def __init__(self, file_list, transform=None, batch_size=1, shuffle=False, num_workers=0): super(CatsDogsDataset, self).__init__(batch_size=batch_size,
wass_metric = args.wass_flag print(wass_metric) x_shape = (channels, img_size, img_size) bce_loss = nn.BCELoss() xe_loss = nn.CrossEntropyLoss() mse_loss = nn.MSELoss() # Initialize generator and discriminator generator = Generator_CNN(latent_dim, n_c, x_shape) encoder = Encoder_CNN(latent_dim, n_c) discriminator = Discriminator_CNN(wass_metric=wass_metric) # Configure data loader transform = transform.Compose([ transform.Resize(size=img_size), transform.Gray(), ]) dataloader = MNIST(train=True, transform=transform).set_attrs(batch_size=batch_size, shuffle=True) testdata = MNIST(train=False, transform=transform).set_attrs(batch_size=batch_size, shuffle=True) (test_imgs, test_labels) = next(iter(testdata)) ge_chain = generator.parameters() for p in encoder.parameters(): ge_chain.append(p) #TODO: weight_decay=decay optimizer_GE = jt.optim.Adam(ge_chain, lr=lr, betas=(b1, b2), weight_decay=0.0)
# Calculate output of image discriminator (PatchGAN) patch_h, patch_w = int(opt.mask_size / 2**3), int(opt.mask_size / 2**3) patch = (1, patch_h, patch_w) # Loss function adversarial_loss = nn.MSELoss() pixelwise_loss = nn.L1Loss() # Initialize generator and discriminator generator = Generator(channels=opt.channels) discriminator = Discriminator(channels=opt.channels) # Dataset loader transforms_ = [ transform.Resize((opt.img_size, opt.img_size), mode=Image.BICUBIC), transform.ImageNormalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ] dataloader = ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transforms_).set_attrs( batch_size=opt.batch_size, shuffle=True, num_workers=opt.n_cpu, ) test_dataloader = ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transforms_, mode="val").set_attrs( batch_size=12, shuffle=True, num_workers=1, )
# Load pretrained models generator.load("saved_models/%s/generator_last.pkl" % (opt.dataset_name)) discriminator.load("saved_models/%s/discriminator_last.pkl" % (opt.dataset_name)) # Optimizers optimizer_G = jt.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) optimizer_D = jt.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) # Configure dataloaders transforms_ = [ transform.Resize(size=(opt.img_height, opt.img_width), mode=Image.BICUBIC), transform.ImageNormalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ] dataloader = ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transforms_).set_attrs( batch_size=opt.batch_size, shuffle=True, num_workers=opt.n_cpu, ) val_dataloader = ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transforms_, mode="val").set_attrs( batch_size=10, shuffle=True, num_workers=1, )
lambda_2 = 100 # ID pixel-wise lambda_3 = 0.1 # KL (encoded translated images) lambda_4 = 100 # Cycle pixel-wise # Optimizers optimizer_G = nn.Adam( E1.parameters() + E2.parameters() + G1.parameters() + G2.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2), ) optimizer_D1 = nn.Adam(D1.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) optimizer_D2 = nn.Adam(D2.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2)) # Image transformations transform_ = [ transform.Resize(int(opt.img_height * 1.12), Image.BICUBIC), transform.RandomCrop((opt.img_height, opt.img_width)), transform.RandomHorizontalFlip(), transform.ImageNormalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ] dataloader = ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transform_, unaligned=True).set_attrs(batch_size=opt.batch_size, shuffle=True, num_workers=opt.n_cpu) val_dataloader = ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transform_, unaligned=True, mode="test").set_attrs(batch_size=5,