def __init__(self, opt):
"""
Parameters:
"""
super(CycleGANModel, self).__init__()
# Stuff
ensure_existance_paths(opt) # Create necessary directories for saving stuff
self.opt = opt
self.gpu_ids = opt.gpu_ids
self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids and torch.cuda.is_available() else torch.device('cpu')
self.lambda_idt = 0.5
self.loss_lambda = 10.0
self.visual_names = ['real_A', 'fake_B', 'rec_A', 'idt_B', 'real_B', 'fake_A', 'rec_B', 'idt_A']
self.name = self.opt.name
# Define models
self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
self.netG_A = networks.to_device(networks.ResNetGenerator(opt.n_input, opt.n_output, opt.forward_mask),self.gpu_ids)
self.netG_B = networks.to_device(networks.ResNetGenerator(opt.n_input, opt.n_output, opt.forward_mask),self.gpu_ids)
self.netD_A = networks.to_device(networks.NLayerDiscriminator(opt.n_output),self.gpu_ids)
self.netD_B = networks.to_device(networks.NLayerDiscriminator(opt.n_output),self.gpu_ids)
# Define losses
self.criterionGAN = networks.GANLoss().to(self.device)
self.criterionCycle = nn.L1Loss().to(self.device)
self.criterionIdt = nn.L1Loss().to(self.device)
self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']
# Define optimizers
self.optimizers = []
self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_D)
# Define schedulers
self.schedulers = [networks.get_optimizer_scheduler(optimizer,opt) for optimizer in self.optimizers]
self.set_train()
self.step = 0
wandb.init(
entity="star-witchers",
project="cycleGAN",
config=self.opt,
name=self.name,
)
def fit(self, data, train_test=False):
# Get the inputs
data = to_device(data, self.args.device)
# If not doing test runs during train times (nearly always).
if not train_test:
left = data['left_image'].squeeze()
return self.G(left)
# Only if also doing tes runs doing train time.
flipped_left = data['flipped_left_image']
left = data['left_image']
test_image_pair = torch.stack((left, flipped_left)).squeeze()
return self.G(test_image_pair)
def __init__(self, opt):
"""
Parameters:
"""
# Stuff
ensure_existance_paths_test(opt) # Create necessary directories for saving stuff
self.opt = opt
self.gpu_ids = opt.gpu_ids
self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids and torch.cuda.is_available() else torch.device('cpu')
self.name = self.opt.name
# Define model
self.netG = networks.to_device(networks.ResNetGenerator(opt.n_input, opt.n_output, opt.forward_mask),self.gpu_ids)
self.netG.eval()
self.load_model(opt.load_model_name)
def set_input(self, data):
self.data = to_device(data, self.device)
self.left = self.data['left_image']
self.right = self.data['right_image']
def predict(self, data):
# Get the inputs
data = to_device(data, self.args.device)
return self.G(data)
