本文整理匯總了Python中datasets.ImageDataset方法的典型用法代碼示例。如果您正苦於以下問題:Python datasets.ImageDataset方法的具體用法?Python datasets.ImageDataset怎麽用?Python datasets.ImageDataset使用的例子?那麽, 這裏精選的方法代碼示例或許可以為您提供幫助。您也可以進一步了解該方法所在類datasets
的用法示例。
在下文中一共展示了datasets.ImageDataset方法的1個代碼示例,這些例子默認根據受歡迎程度排序。您可以為喜歡或者感覺有用的代碼點讚,您的評價將有助於係統推薦出更棒的Python代碼示例。
示例1: main
# 需要導入模塊: import datasets [as 別名]
# 或者: from datasets import ImageDataset [as 別名]
def main():
device = torch.device("cuda:0" if FLAGS.cuda else "cpu")
if FLAGS.train:
print('Loading data...\n')
transform = [transforms.Resize(int(FLAGS.img_size*1.12), Image.BICUBIC),
transforms.RandomCrop((FLAGS.img_size, FLAGS.img_size)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))]
dataloader = DataLoader(ImageDataset(os.path.join(FLAGS.data_dir, FLAGS.dataset),
transform=transform, unaligned=True, mode='train'),
batch_size=FLAGS.batch_size, shuffle=True, num_workers=2)
test_dataloader = DataLoader(ImageDataset(os.path.join(FLAGS.data_dir, FLAGS.dataset),
transform=transform, unaligned=True, mode='test'),
batch_size=FLAGS.test_batch_size, shuffle=True, num_workers=2)
print('Creating model...\n')
model = Model(FLAGS.model, device, dataloader, test_dataloader, FLAGS.channels, FLAGS.img_size, FLAGS.num_blocks)
model.create_optim(FLAGS.lr)
# Train
model.train(FLAGS.epochs, FLAGS.log_interval, FLAGS.out_dir, True)
model.save_to('')
else:
model = Model(FLAGS.model, device, None, test_dataloader, FLAGS.channels, FLAGS.img_size, FLAGS.num_blocks)
model.load_from(FLAGS.out_dir)
model.eval(mode=1, batch_size=FLAGS.batch_size)
開發者ID:PacktPublishing,項目名稱:Hands-On-Generative-Adversarial-Networks-with-PyTorch-1.x,代碼行數:31,代碼來源:main.py