def __init__(self,
height_lr=24,
width_lr=24,
channels=3,
upscaling_factor=4,
gen_lr=1e-4,
dis_lr=1e-4,
loss_weights=[0.006, 1e-4],
training_mode=True,
colorspace='RGB'):
"""
:param int height_lr: Height of low-resolution images
:param int width_lr: Width of low-resolution images
:param int channels: Image channels
:param int upscaling_factor: Up-scaling factor
:param int gen_lr: Learning rate of generator
:param int dis_lr: Learning rate of discriminator
"""
# Low-resolution image dimensions
self.height_lr = height_lr
self.width_lr = width_lr
self.training_mode = training_mode
# High-resolution image dimensions
if upscaling_factor not in [2, 4, 8]:
raise ValueError(
'Upscaling factor must be either 2, 4, or 8. You chose {}'.
format(upscaling_factor))
self.upscaling_factor = upscaling_factor
self.height_hr = int(self.height_lr * self.upscaling_factor)
self.width_hr = int(self.width_lr * self.upscaling_factor)
# Low-resolution and high-resolution shapes
self.channels = channels
self.colorspace = colorspace
self.shape_lr = (self.height_lr, self.width_lr, self.channels)
self.shape_hr = (self.height_hr, self.width_hr, self.channels)
# Learning rates
self.gen_lr = gen_lr
self.dis_lr = dis_lr
# Gan setup settings
self.loss_weights = loss_weights
self.VGGLoss = VGGLoss(self.shape_hr)
self.gen_loss = 'mse'
self.content_loss = self.VGGLoss.content_loss # self.VGGLoss.euclidean_content_loss
self.adversarial_loss = 'binary_crossentropy'
# Build & compile the generator network
self.generator = self.build_generator()
self.compile_generator(self.generator)
# If training, build rest of GAN network
if training_mode:
self.discriminator = self.build_discriminator()
self.compile_discriminator(self.discriminator)
self.srgan = self.build_srgan()
self.compile_srgan(self.srgan)
# loss 関数の設定
#======================================================================
# Adv loss
if (args.gan_type == "vanilla"):
loss_gan_fn = VanillaGANLoss(device, w_sigmoid_D=False)
elif (args.gan_type == "lsgan"):
loss_gan_fn = LSGANLoss(device)
else:
raise NotImplementedError('gan_type %s not implemented' %
args.gan_type)
# feature maching loss
loss_feat_fn = FeatureMatchingLoss(device, n_dis=args.n_dis, n_layers_D=3)
# vgg perceptual loss
loss_vgg_fn = VGGLoss(device)
#======================================================================
# モデルの学習処理
#======================================================================
# 学習中の生成画像の履歴
fake_images_historys = []
print("Starting Training Loop...")
total_steps = 0
iterations = 0 # 学習処理のイテレーション回数
n_print = 1
#-----------------------------
# エポック数分トレーニング
#-----------------------------
for epoch in tqdm(range(args.n_epoches), desc="Epoches"):
def __init__(self,
height_lr=64,
width_lr=64,
channels=3,
upscaling_factor=2,
gen_lr=1e-3,
dis_lr=1e-4,
ra_lr=1e-4,
loss_weights=[1e-3, 5e-3, 1e-2],
training_mode=True,
colorspace='RGB'):
"""
height_lr: height of the lr image
width_lr: width of the lr image
channels: number of channel of the image
upscaling_factor= factor upscaling
lr = learning rate
training_mode: True or False
colorspace: 'RGB' or 'YCbCr'
"""
# Low-resolution image dimensions
self.height_lr = height_lr
self.width_lr = width_lr
self.training_mode = training_mode
# High-resolution image dimensions
if upscaling_factor not in [2, 3, 4, 8]:
raise ValueError(
'Upscaling factor must be either 2, 4, or 8. You chose {}'.
format(upscaling_factor))
self.upscaling_factor = upscaling_factor
self.height_hr = int(self.height_lr * self.upscaling_factor)
self.width_hr = int(self.width_lr * self.upscaling_factor)
# Low-resolution and high-resolution shapes
self.channels = channels
self.colorspace = colorspace
self.shape_lr = (self.height_lr, self.width_lr, self.channels)
self.shape_hr = (self.height_hr, self.width_hr, self.channels)
# Learning rates
self.gen_lr = gen_lr
self.dis_lr = dis_lr
self.ra_lr = ra_lr
# Gan setup settings
self.loss_weights = loss_weights
self.VGGLoss = VGGLoss(self.shape_hr)
self.gen_loss = euclidean #euclidean, charbonnier, mae,mse,L2Loss,L1Loss,L1L2EucliLoss
self.content_loss = L1Loss #self.VGGLoss.content_loss
self.adversarial_loss = binary_crossentropy
self.ra_loss = binary_crossentropy
# Build & compile the generator network
self.generator = self.build_generator()
self.compile_generator(self.generator)
# If training, build rest of GAN network
if training_mode:
self.discriminator = self.build_discriminator()
self.compile_discriminator(self.discriminator)
self.ra_discriminator = self.build_ra_discriminator()
self.compile_ra_discriminator(self.ra_discriminator)
self.rtvsrgan = self.build_rtvsrgan()
self.compile_rtvsrgan(self.rtvsrgan)
def __init__(self, upscale_factor=4, generator_weights=None, **kwargs):
super(SRGan, self).__init__(**kwargs)
self.generator = Generator(weights=generator_weights)
self.discriminator = Discriminator()
self.vgg_loss = VGGLoss()
torch.save(
{
'epoch': epoch,
'gen': net.state_dict(),
'gen_opt': hsid_optimizer.state_dict()
}, os.path.join('./checkpoints', "model_latest.pth"))
tb_writer.close()
from model_refactored import HSIDDenseNetTwoStage
from Discriminator_copy import DiscriminatorABC
adv_criterion = nn.BCEWithLogitsLoss()
recon_criterion = nn.L1Loss()
lambda_recon = 10
vgg_loss = VGGLoss().cuda(device=DEVICE)
def train_model_residual_lowlight_twostage_gan_best():
start_epoch = 1
device = DEVICE
#准备数据
train_set = HsiCubicTrainDataset('./data/train_lowlight/')
print('total training example:', len(train_set))
train_loader = DataLoader(dataset=train_set,
batch_size=BATCH_SIZE,
shuffle=True)
def build_model(self):
# Generator = Scene Graph Model
kwargs = {
'vocab': self.vocab,
'image_size': self.image_size,
'embedding_dim': self.embedding_dim,
'gconv_dim': self.gconv_dim,
'gconv_hidden_dim': self.gconv_hidden_dim,
'gconv_num_layers': self.gconv_num_layers,
'mlp_normalization': self.mlp_normalization,
'normalization': self.normalization,
'activation': self.activation,
'mask_size': self.mask_size,
'layout_noise_dim': self.layout_noise_dim,
'netG': self.netG,
'ngf': self.ngf,
'n_downsample_global': self.n_downsample_global,
'n_blocks_global': self.n_blocks_global,
'n_blocks_local': self.n_blocks_local,
'n_local_enhancers': self.n_local_enhancers,
}
self.generator = Sg2ImModel(**kwargs)
# Discriminators
# OBJ Discriminator
# self.obj_discriminator = AcCropDiscriminator(vocab=self.vocab,
# arch=self.d_obj_arch,
# normalization=self.d_normalization,
# activation=self.d_activation,
# padding=self.d_padding,
# object_size=self.crop_size)
self.obj_discriminator = ref_network.define_obj_D(
vocab=self.vocab,
input_nc=3,
crop_size=self.crop_size,
ndf=self.ndf,
n_layers_D=self.n_layers_D_obj,
norm=self.normalization,
num_D=self.num_D_obj,
getIntermFeat=not self.no_objFeat_loss)
# IMG Discriminator
self.img_discriminator = ref_network.define_img_D(
input_nc=3,
ndf=self.ndf,
n_layers_D=self.n_layers_D_img,
norm=self.normalization,
num_D=self.num_D_img,
getIntermFeat=not self.no_imgFeat_loss)
# Loss
self.gan_g_loss, self.gan_d_loss = get_gan_losses(self.gan_loss_type)
self.vgg_loss = VGGLoss()
# Weights
# feature matching loss weights
self.img_feat_weights = 4.0 / (self.n_layers_D_img + 1)
self.img_D_weights = 1.0 / self.num_D_img
self.obj_feat_weights = 4.0 / (self.n_layers_D_obj + 1)
self.obj_D_weights = 1.0 / self.num_D_obj
# Mode of Model
if self.init_epochs >= self.eval_mode_after:
print("Setting Generator to Eval Mode")
self.generator.eval()
else:
print("Setting Generator to Train Mode")
self.generator.train()
self.obj_discriminator.train()
self.img_discriminator.train()
# Optimizers
self.gen_optimizer = torch.optim.Adam(self.generator.parameters(),
lr=self.learning_rate)
self.dis_obj_optimizer = torch.optim.Adam(
self.obj_discriminator.parameters(), lr=self.learning_rate)
self.dis_img_optimizer = torch.optim.Adam(
self.img_discriminator.parameters(), lr=self.learning_rate)
# Others
self.ac_ind = 1
self.obj_ind = 0
if not self.no_objFeat_loss:
self.ac_ind = 2
self.obj_ind = 1
# Print networks
self.print_network(self.generator, 'Generator')
self.print_network(self.obj_discriminator, 'Object Discriminator')
self.print_network(self.img_discriminator, 'Image Discriminator')
if torch.cuda.is_available():
self.generator.cuda()
self.obj_discriminator.cuda()
self.img_discriminator.cuda()