import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.utils as vutils
from itertools import chain
from utils.dataset import DATASET
from utils.ImagePool import ImagePool
from model.Discriminator import Discriminator
from model.Generator import Generator
########## DATASET ###########
datasetA = DATASET(os.path.join("facades/train/",'A'),256,256,1)
datasetB = DATASET(os.path.join("facades/train/",'B'),256,256,1)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
os.path.join("facades/",'train'),
transforms.Compose([
transforms.RandomResizedCrop(256),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
loader_A = torch.utils.data.DataLoader(dataset=datasetA,
batch_size=1,
shuffle=True,
num_workers=2)
try:
os.makedirs(opt.outf)
except OSError:
pass
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
torch.cuda.manual_seed_all(opt.manualSeed)
cudnn.benchmark = True
########## DATASET ###########
datasetA = DATASET(os.path.join(opt.dataPath, 'A'), opt.loadSize, opt.fineSize,
opt.flip)
datasetB = DATASET(os.path.join(opt.dataPath, 'B'), opt.loadSize, opt.fineSize,
opt.flip)
loader_A = torch.utils.data.DataLoader(dataset=datasetA,
batch_size=opt.batchSize,
shuffle=True,
num_workers=2)
loaderA = iter(loader_A)
loader_B = torch.utils.data.DataLoader(dataset=datasetB,
batch_size=opt.batchSize,
shuffle=True,
num_workers=2)
loaderB = iter(loader_B)
########### MODEL ###########
ndf = opt.ndf
ngf = opt.ngf
os.makedirs(opt.dataroot)
except OSError:
pass
try:
os.makedirs(opt.outf)
except OSError:
pass
device = torch.device("cuda:0" if torch.cuda.is_available else "cpu")
# dataset = dataloader(opt.dataroot, opt.datasetA, opt.datasetB, bs=opt.batchSize, if_generate=False)
# datasetA = DATASET('/home/szk/PycharmProjects/open-reid/examples/data/cuhk03/images', 143, 128, 1)
# datasetB = DATASET('/home/szk/PycharmProjects/open-reid/examples/data/viper/images', 143, 128, 1)
datasetA = DATASET(
'/home/szk/PycharmProjects/pytorch-CycleGAN-and-pix2pix/datasets/horse2zebra/trainA',
143, 128, 1)
datasetB = DATASET(
'/home/szk/PycharmProjects/pytorch-CycleGAN-and-pix2pix/datasets/horse2zebra/trainB',
143, 128, 1)
loader_A = torch.utils.data.DataLoader(dataset=datasetA,
batch_size=opt.batchSize,
shuffle=True)
loaderA = iter(loader_A)
loader_B = torch.utils.data.DataLoader(dataset=datasetB,
batch_size=opt.batchSize,
shuffle=True)
loaderB = iter(loader_B)
writer = SummaryWriter()