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,
