def test_random_crop(self):
height = random.randint(10, 32) * 2
width = random.randint(10, 32) * 2
oheight = random.randint(5, (height - 2) / 2) * 2
owidth = random.randint(5, (width - 2) / 2) * 2
img = np.ones((height, width, 3))
result = transform.Compose([
transform.ToPILImage(),
transform.RandomCrop((oheight, owidth)),
transform.ToTensor(),
])(img)
self.assertEqual(result.shape[1], oheight)
self.assertEqual(result.shape[2], owidth)
result = transform.Compose([
transform.ToPILImage(),
transform.RandomCrop((oheight, owidth)),
transform.ToTensor(),
])(img)
self.assertEqual(result.shape[1], oheight)
self.assertEqual(result.shape[2], owidth)
result = transform.Compose([
transform.ToPILImage(),
transform.RandomCrop((height, width)),
transform.ToTensor()
])(img)
self.assertEqual(result.shape[1], height)
self.assertEqual(result.shape[2], width)
self.assertTrue(np.allclose(img, result.transpose(1, 2, 0)))
with self.assertRaises(AssertionError):
result = transform.Compose([
transform.ToPILImage(),
transform.RandomCrop((height + 1, width + 1)),
transform.ToTensor(),
])(img)
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)
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,
shuffle=True,
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)),
]