def __init__(self, config, mode):
super().__init__()
self.config = config
self.mode = mode
self.h2l_G = HighToLowGenerator()
self.l2h_G = LowToHighGenerator()
self.is_cuda = torch.cuda.is_available()
if self.is_cuda and not self.config.cuda:
print(
"WARNING: You have a CUDA device, so you should probably enable CUDA"
)
self.cuda = self.is_cuda & self.config.cuda
self.manual_seed = random.randint(1, 10000)
print('seed:{}'.format(self.manual_seed))
random.seed(self.manual_seed)
self.test_file = self.config.output_path
if not os.path.exists(self.test_file):
os.makedirs(self.test_file)
if self.cuda:
self.device = torch.device("cuda")
torch.cuda.set_device(self.config.gpu_device)
torch.cuda.manual_seed_all(self.manual_seed)
print("Program will run on *****GPU-CUDA***** ")
print_cuda_statistics()
else:
self.device = torch.device("cpu")
torch.manual_seed(self.manual_seed)
print("Program will run on *****CPU***** ")
self.l2h_G = self.l2h_G.to(self.device)
self.h2l_G = self.h2l_G.to(self.device)
def __init__(self, config):
super().__init__(config)
# define models (generator and discriminator)
self.h2l_G = HighToLowGenerator()
self.h2l_D = HighToLowDiscriminator()
self.l2h_G = LowToHighGenerator()
self.l2h_D = LowToHighDiscriminator()
# define loss
#self.loss = GANLoss()
#self.loss = HingeEmbeddingLoss()
self.criterion_GAN = torch.nn.BCEWithLogitsLoss()
self.criterion_MSE = MSELoss()
# define optimizers for both generator and discriminator
self.l2h_optimG = torch.optim.Adam(self.l2h_G.parameters(),
lr=self.config.learning_rate,
betas=(self.config.beta1,
self.config.beta2))
self.l2h_optimD = torch.optim.Adam(self.l2h_D.parameters(),
lr=self.config.learning_rate,
betas=(self.config.beta1,
self.config.beta2))
self.h2l_optimG = torch.optim.Adam(self.h2l_G.parameters(),
lr=self.config.learning_rate,
betas=(self.config.beta1,
self.config.beta2))
self.h2l_optimD = torch.optim.Adam(self.h2l_D.parameters(),
lr=self.config.learning_rate,
betas=(self.config.beta1,
self.config.beta2))
# initialize counter
self.current_epoch = 0
self.current_iteration = 0
self.best_valid_mean_iou = 0
self.real_label = 1
self.fake_label = -1
# set cuda flag
self.is_cuda = torch.cuda.is_available()
if self.is_cuda and not self.config.cuda:
self.logger.info(
"WARNING: You have a CUDA device, so you should probably enable CUDA"
)
self.cuda = self.is_cuda & self.config.cuda
# set the manual seed for torch
self.manual_seed = random.randint(1, 10000)
self.logger.info('seed:{}'.format(self.manual_seed))
random.seed(self.manual_seed)
self.test_file = self.config.output_path
if not os.path.exists(self.test_file):
os.makedirs(self.test_file)
if self.cuda:
self.device = torch.device("cuda")
torch.cuda.set_device(self.config.gpu_device)
torch.cuda.manual_seed_all(self.manual_seed)
self.logger.info("Program will run on *****GPU-CUDA***** ")
print_cuda_statistics()
else:
self.device = torch.device("cpu")
torch.manual_seed(self.manual_seed)
self.logger.info("Program will run on *****CPU***** ")
self.l2h_G = self.l2h_G.to(self.device)
self.l2h_D = self.l2h_D.to(self.device)
self.h2l_G = self.h2l_G.to(self.device)
self.h2l_D = self.h2l_D.to(self.device)
self.criterion_GAN = self.criterion_GAN.to(self.device)
self.criterion_MSE = self.criterion_MSE.to(self.device)
# Summary Writer
self.summary_writer_l2h = SummaryWriter(
log_dir=self.config.summary_dir_l2h, comment='Low-To-High GAN')
self.summary_writer_h2l = SummaryWriter(
log_dir=self.config.summary_dir_h2l, comment='High-To-Low GAN')
