from torch.utils.data import DataLoader
from networks import Discriminator, Generator, Loss
from option import TrainOption, TestOption
from pipeline import CustomDataset
from utils import Manager, update_lr, weights_init
import datetime
torch.backends.cudnn.benchmark = True
opt = TrainOption().parse()
os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu_ids)
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu:0')
lr = opt.lr
dataset = CustomDataset(opt)
data_loader = DataLoader(dataset=dataset,
batch_size=opt.batch_size,
num_workers=opt.n_workers,
shuffle=opt.shuffle)
test_opt = TestOption().parse()
test_dataset = CustomDataset(test_opt)
test_data_loader = DataLoader(dataset=test_dataset,
batch_size=test_opt.batch_size,
num_workers=test_opt.n_workers)
G = Generator(opt).apply(weights_init).to(DEVICE)
D = Discriminator(opt).apply(weights_init).to(DEVICE)
criterion = Loss(opt)
A_optim = torch.optim.Adam(A.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2), eps=opt.eps)
current_step = 0
epoch_decay = opt.epoch_decay
lr = opt.lr
n_critics = opt.n_critics
n_epochs = opt.n_epochs
n_epochs_per_lod = opt.n_epochs_per_lod
n_iter_per_lod = n_epochs_per_lod * opt.n_data
nb_transition = n_iter_per_lod / 2
package = {}
start_time = datetime.datetime.now()
if opt.progression:
for level in range(opt.n_downsample + 1): # 0 1 2 3 4 5
level_in = level
dataset = CustomDataset(opt, level=level)
data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=opt.batch_size,
num_workers=opt.n_workers, shuffle=opt.shuffle)
for epoch in range(n_epochs_per_lod):
package.update({'Epoch': epoch + 1})
for _, data_dict in enumerate(data_loader):
current_step += 1
level_in += 1 / nb_transition
level_in = np.clip(level_in, level, level + 1.0)
for k, v in data_dict.items():
data_dict.update({k: v.to(device)})
package.update(criterion(A, G, data_dict, level=level, level_in=level_in))
import torch
import glob
from pipeline import CustomDataset
from network import Generator
from option import TestOption
from utils.Manger import save_image
import datetime
if __name__ == '__main__':
start_time = datetime.datetime.now()
opt = TestOption().parse()
USE_CUDA = True if torch.cuda.is_available() else False
dataset = CustomDataset(opt)
data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=opt.batch_size,
shuffle=opt.is_train, num_workers=opt.n_workers)
G = Generator(opt)
list_G = sorted(glob.glob(os.path.join(opt.model_dir, '*_G.pt')))
for path_G in list_G:
G.load_state_dict(torch.load(path_G))
G = G.cuda() if USE_CUDA else G
for i, input in enumerate(data_loader):
input = input.cuda() if USE_CUDA else input
fake = G(input)
save_image(fake.detach(), os.path.join(opt.image_dir, 'Result_{}'.format(i)))
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from models import Discriminator, Generator
from pipeline import CustomDataset
DEVICE = torch.device("gpu:0" if torch.cuda.is_available() else "cpu:0")
DIR_IMAGE = './pix2pix/checkpoints/IMAGE'
DIR_MODEL = './pix2pix/checkpoints/MODEL'
EPOCHS = 10
EPOCH_SAVE = 5
ITER_DISPLAY = 10
os.makedirs(DIR_IMAGE) if not os.path.isdir(DIR_IMAGE) else None
os.makedirs(DIR_MODEL) if not os.path.isdir(DIR_MODEL) else None
dataset = CustomDataset(root='./datasets/Noel', crop_size=128, flip=True)
data_loader = DataLoader(dataset=dataset,
batch_size=1,
shuffle=True,
num_workers=1)
D = Discriminator()
G = Generator()
GAN_Loss = nn.BCELoss()
L1_Loss = nn.L1Loss()
D_optim = torch.optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
G_optim = torch.optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
total_iter = 0
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from Pix2Pix_models import Discriminator, Generator
from pipeline import CustomDataset
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
DIR_IMAGE = 'pix2pix\\checkpoint\\IMAGE\\'
DIR_MODEL = 'pix2pix\\checkpoint\\MODEL\\'
EPOCHS = 10
EPOCH_SAVE = 5
ITER_DISPLAY = 10
os.makedirs(DIR_IMAGE) if not os.path.isdir(DIR_IMAGE) else None
os.makedirs(DIR_MODEL) if not os.path.isdir(DIR_MODEL) else None
dataset = CustomDataset(root='dataset\\images\\', crop_size=128, flip=True)
data_loader = DataLoader(dataset=dataset,
batch_size=1,
shuffle=True,
num_workers=1)
D = Discriminator()
G = Generator()
GAN_Loss = nn.BCELoss()
L1_Loss = nn.L1Loss()
D_optim = torch.optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
G_optim = torch.optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
total_iter = 0
iter_report = opt.iter_report
iter_val = opt.iter_val
n_epochs = opt.n_epochs
lambda_cycle = opt.lambda_cycle
lr = opt.lr
n_buffer_images = opt.n_buffer_images
num_workers = opt.n_workers
val_during_training = opt.val_during_training
if val_during_training:
from options import TestOption
opt_test = TestOption().parse()
dir_image_test = opt.dir_image_test
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu:0')
dataset = CustomDataset(opt, 'train')
data_loader = DataLoader(dataset=dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
D_A = Discriminator(opt).apply(weights_init).to(device)
D_B = Discriminator(opt).apply(weights_init).to(device)
G_A = Generator(opt).apply(weights_init).to(device)
G_B = Generator(opt).apply(weights_init).to(device)
print(D_A)
print(G_A)
optim_D = torch.optim.Adam(list(D_A.parameters()) + list(D_B.parameters()),
lr=lr,
betas=(beta_1, beta_2))