parser.add_argument("--n_cpu", type=int, default=0, help="number of cpu threads to use during batch generation") parser.add_argument("--nz", type=int, default=10, help="dimensionality of the latent space") parser.add_argument("--imageSize", type=int, default=28, help="size of each image dimension") parser.add_argument("--sample_interval", type=int, default=100, help="interval between image sampling") opt = parser.parse_args() print(opt) os.makedirs(opt.experiment, exist_ok=True) ## device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(device) ## load dataset os.makedirs(os.path.join(opt.dataroot, opt.dataset), exist_ok=True) trans = transforms.Compose([transforms.ToTensor()]) dataset = load_dataset(opt.dataroot, opt.dataset, opt.imageSize, trans, train=False) assert dataset dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize, shuffle=True, num_workers=opt.n_cpu) opt.dataSize = len(dataset) ## parameters opt.n_classes = len(dataset.classes) ## load pretrained model enc = Encoder(opt.imageSize**2, opt.nz, opt.n_classes).to(device) dec = Decoder(opt.imageSize**2, opt.nz, opt.n_classes).to(device) enc.load_state_dict(torch.load(opt.enc_pth)) dec.load_state_dict(torch.load(opt.dec_pth)) print("Pretrained models have been loaded.")
# opt.seed = 42 # torch.manual_seed(opt.seed) # np.random.seed(opt.seed) ## cudnn torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = True ## device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(device) ## dataset train_dataset = load_dataset(opt.dataroot, opt.dataset, opt.imageSize, trans=None, train=True) train_dataloader = DataLoader(train_dataset, batch_size=opt.batchSize, shuffle=True) ## model init def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv') != -1: m.weight.data.normal_(0.0, 0.02) elif classname.find('BatchNorm') != -1: m.weight.data.normal_(1.0, 0.02) m.bias.data.fill_(0)
## random seed # opt.seed = 42 # torch.manual_seed(opt.seed) # np.random.seed(opt.seed) ## cudnn torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = True ## device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(device) ## dataset test_dataset = load_dataset(opt.dataroot, opt.dataset, opt.size, trans=None, train=False) test_dataloader = DataLoader(test_dataset, batch_size=opt.batchSize, shuffle=False) opt.dataSize = test_dataset.__len__() ## model gen = NetG(opt).to(device) disc = NetD(opt).to(device) assert gen.load_state_dict(torch.load(opt.gen_pth)) assert disc.load_state_dict(torch.load(opt.disc_pth)) print("Pretrained models have been loaded.") ## record results writer = SummaryWriter("../runs{0}".format(opt.experiment[1:]), comment=opt.experiment[1:]) ## def def splitImage(img, size):