dset = VideoFolderDataset(args.path,
transform,
mode='image',
cache=args.cache)
elif args.dataset == 'imagefolder':
transform = transforms.Compose([
transforms.RandomHorizontalFlip(p=0.5 if args.flip else 0),
transforms.Resize(args.size, Image.LANCZOS),
transforms.CenterCrop(args.size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
dset = datasets.ImageFolder(args.path, transform=transform)
else:
dset = get_image_dataset(args,
args.dataset,
args.path,
train=args.eval_type == 'train')
# args.n_sample = min(args.n_sample, len(dset))
indices = torch.randperm(len(dset))[:args.n_sample]
dset = Subset(dset, indices)
loader = DataLoader(dset,
batch_size=args.batch,
num_workers=4,
shuffle=True)
features = extract_features(loader, inception, device).numpy()
# features = features[: args.n_sample]
print(f"extracted {features.shape[0]} features")
d_optim.load_state_dict(ckpt["d_optim"])
if args.distributed:
generator = nn.parallel.DistributedDataParallel(
generator,
device_ids=[args.local_rank],
output_device=args.local_rank,
broadcast_buffers=False,
)
discriminator = nn.parallel.DistributedDataParallel(
discriminator,
device_ids=[args.local_rank],
output_device=args.local_rank,
broadcast_buffers=False,
)
dataset = get_image_dataset(args, args.dataset, args.path, train=True)
loader = data.DataLoader(
dataset,
batch_size=args.batch,
sampler=data_sampler(dataset, shuffle=True, distributed=args.distributed),
drop_last=True,
)
if get_rank() == 0 and wandb is not None and args.wandb:
wandb.init(project=args.name)
util.print_models([generator, discriminator], args)
train(args, loader, generator, discriminator, g_optim, d_optim, g_ema, device)
