torch.backends.cudnn.benchmark = True
transform = transforms.Compose(
[
transforms.Resize(args.size),
transforms.CenterCrop(args.size),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
]
)
# TODO use a little set for sanity check
# _, loader, _ = iPERLoader(data_root=args.path, batch=args.batch_size, transform=transform).data_load()
_, _, loader = iPERLoader(data_root=args.path, batch=args.batch_size, transform=transform).data_load()
# model for image
model_img = AppVQVAE().to(device)
model_img = nn.DataParallel(model_img).to(device)
if is_load_model_img is True:
print('Loading model_img ...', end='')
model_img.load_state_dict(torch.load(args.model_img_path))
model_img.eval()
print('Done')
else:
print('model_img Initialized.')
optimizer_img = optim.Adam(model_img.parameters(), lr=args.lr)
# model for condition
model_cond = VQVAE().to(device)
model_cond = nn.DataParallel(model_cond).cuda()
if is_load_model_cond is True:
print('Loading model_cond ...', end='')
torch.backends.cudnn.benchmark = True
transform = transforms.Compose([
transforms.Resize(args.size),
transforms.CenterCrop(args.size),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])
# TODO use a little set for sanity check
# _, loader, _ = iPERLoader(data_root=args.path, batch=args.batch_size, transform=transform).data_load()
_, _, loader = iPERLoader(data_root=args.path,
batch=args.batch_size,
transform=transform).data_load()
# model for image
model_img = AppVQVAE().to(device)
model_img = nn.DataParallel(model_img).to(device)
if is_load_model_img is True:
print('Loading model_img ...', end='')
model_img.load_state_dict(torch.load(args.model_img_path))
model_img.eval()
print('Done')
else:
print('model_img Initialized.')
optimizer_img = optim.Adam(model_img.parameters(), lr=args.lr)
# model for condition
model_cond = VQVAE().to(device)
model_cond = nn.DataParallel(model_cond).cuda()
if is_load_model_cond is True:
print('Loading model_cond ...', end='')