def __init__(self,
width,
height,
ichan,
ochan,
l1_weight=100.,
lr=0.0002,
beta1=0.5):
# def __init__(self, width, height, ichan, ochan, l1_weight=1000., lr=0.0002, beta1=0.5): # I enlarge the l1_weight to see what will happen
"""
width: image width in pixel.
height: image height in pixel.
ichan: number of channels used by input images.
ochan: number of channels used by output images.
l1_weight: L1 loss weight.
lr: learning rate for ADAM optimizer.
beta1: beta1 parameter for ADAM optimizer.
"""
self._is_training = tf.placeholder(tf.bool)
self._g_inputs = tf.placeholder(tf.float32,
[None, width, height, ichan])
self._d_inputs_a = tf.placeholder(tf.float32,
[None, width, height, ichan])
self._d_inputs_b = tf.placeholder(tf.float32,
[None, width, height, ochan])
self._g = Generator(self._g_inputs, self._is_training, ochan)
self._real_d = Discriminator(
tf.concat([self._d_inputs_a, self._d_inputs_b], axis=3),
self._is_training)
self._fake_d = Discriminator(tf.concat(
[self._d_inputs_a, self._g._decoder['cl9']['fmap']], axis=3),
self._is_training,
reuse=True)
self._g_loss = -tf.reduce_mean(
tf.log(self._fake_d._discriminator['l5']['fmap'])
) + l1_weight * tf.reduce_mean(
tf.abs(self._d_inputs_b - self._g._decoder['cl9']['fmap']))
self._d_loss = -tf.reduce_mean(
tf.log(self._real_d._discriminator['l5']['fmap']) +
tf.log(1.0 - self._fake_d._discriminator['l5']['fmap']))
# self._g_loss = tf.reduce_mean(tf.square(self._d_inputs_b - self._g._decoder['cl9']['fmap'])) # set g_loss = mse loss !!!!!
# self._d_loss = -tf.reduce_mean(tf.log(self._real_d._discriminator['l5']['fmap']) + tf.log(1.0 - self._fake_d._discriminator['l5']['fmap']))
g_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope='G')
with tf.control_dependencies(g_update_ops):
self._g_train_step = tf.train.AdamOptimizer(
lr, beta1=beta1).minimize(self._g_loss,
var_list=tf.get_collection(
tf.GraphKeys.TRAINABLE_VARIABLES,
scope='G'))
d_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope='D')
with tf.control_dependencies(d_update_ops):
self._d_train_step = tf.train.AdamOptimizer(
lr, beta1=beta1).minimize(self._d_loss,
var_list=tf.get_collection(
tf.GraphKeys.TRAINABLE_VARIABLES,
scope='D'))
def build_model(self):
self.Encoder = Encoder(self.in_channels, self.storing_channels,
self.nf).to(self.device)
self.Decoder = Decoder(self.storing_channels, self.in_channels,
self.nf).to(self.device)
self.Disciminator = Discriminator().to(self.device)
self.load_model()
def __init__(self, cfg):
super(RefSRSolver, self).__init__(cfg)
self.srntt = SRNTT(cfg['model']['n_resblocks'],
cfg['schedule']['use_weights'],
cfg['schedule']['concat']).cuda()
# self.discriminator = None
self.discriminator = Discriminator(cfg['data']['input_size']).cuda()
# self.vgg = None
self.vgg = VGG19(cfg['model']['final_layer'],
cfg['model']['prev_layer'], True).cuda()
params = list(self.srntt.texture_transfer.parameters()) + list(self.srntt.texture_fusion_medium.parameters()) +\
list(self.srntt.texture_fusion_large.parameters()) + list(self.srntt.srntt_out.parameters())
self.init_epoch = self.cfg['schedule']['init_epoch']
self.num_epochs = self.cfg['schedule']['num_epochs']
self.optimizer_init = torch.optim.Adam(params,
lr=cfg['schedule']['lr'])
self.optimizer = torch.optim.lr_scheduler.MultiStepLR(
torch.optim.Adam(params, lr=cfg['schedule']['lr']),
[self.num_epochs // 2], 0.1)
self.optimizer_d = torch.optim.lr_scheduler.MultiStepLR(
torch.optim.Adam(self.discriminator.parameters(),
lr=cfg['schedule']['lr']), [self.num_epochs // 2],
0.1)
self.reconst_loss = nn.L1Loss()
self.bp_loss = BackProjectionLoss()
self.texture_loss = TextureLoss(self.cfg['schedule']['use_weights'],
80)
self.adv_loss = AdvLoss(self.cfg['schedule']['is_WGAN_GP'])
self.loss_weights = self.cfg['schedule']['loss_weights']
def test():
# load std models
# policy_log_std = torch.load('./model_pkl/policy_net_action_std_model_1.pkl')
# transition_log_std = torch.load('./model_pkl/transition_net_state_std_model_1.pkl')
# define actor/critic/discriminator net and optimizer
policy = Policy(discrete_action_sections, discrete_state_sections)
value = Value()
discriminator = Discriminator()
discriminator_criterion = nn.BCELoss()
# load expert data
dataset = ExpertDataSet(args.data_set_path)
data_loader = data.DataLoader(dataset=dataset,
batch_size=1,
shuffle=False,
num_workers=0)
# load net models
discriminator.load_state_dict(
torch.load('./model_pkl/Discriminator_model_2.pkl'))
policy.transition_net.load_state_dict(
torch.load('./model_pkl/Transition_model_2.pkl'))
policy.policy_net.load_state_dict(
torch.load('./model_pkl/Policy_model_2.pkl'))
value.load_state_dict(torch.load('./model_pkl/Value_model_2.pkl'))
discrete_state_loss_list = []
continous_state_loss_list = []
action_loss_list = []
cnt = 0
for expert_state_batch, expert_action_batch, expert_next_state in data_loader:
cnt += 1
expert_state_action = torch.cat(
(expert_state_batch, expert_action_batch),
dim=-1).type(FloatTensor)
next_discrete_state, next_continuous_state, _ = policy.get_transition_net_state(
expert_state_action)
gen_next_state = torch.cat(
(next_discrete_state.to(device), next_continuous_state.to(device)),
dim=-1)
loss_func = torch.nn.MSELoss()
continous_state_loss = loss_func(gen_next_state[:, 132:],
expert_next_state[:, 132:])
discrete_state_loss = hamming_loss(
gen_next_state[:, :132],
expert_next_state[:, :132].type(torch.LongTensor))
discrete_action, continuous_action, _ = policy.get_policy_net_action(
expert_state_batch.type(FloatTensor))
gen_action = torch.FloatTensor(continuous_action)
loss_func = torch.nn.MSELoss()
action_loss = loss_func(gen_action, expert_action_batch)
discrete_state_loss_list.append(discrete_state_loss)
continous_state_loss_list.append(continous_state_loss.item())
action_loss_list.append(action_loss)
print(sum(discrete_state_loss_list) / cnt)
print(sum(continous_state_loss_list) / cnt)
def __init__(self):
super(Network, self).__init__()
self.feature_extractor = resnet50() # Already pretrained
# self.feature_extractor = resnet50(pretrained_path=None)
self.selector = Selector()
self.dis = Discriminator()
self.optmzr_select = Adam(self.selector.parameters(), lr=1e-3)
self.optmzr_dis = Adam(self.dis.parameters(), lr=1e-3)
def restore(self):
print('Restoring model...')
generator = Generator(self.resolution, self.batch_size)
generator.restore()
discriminator = Discriminator(self.resolution)
discriminator.restore()
self.restored = True
self.initialize(generator, discriminator)
def __init__(self, model_dir, g_optimizer, d_optimizer, lr, warmup,
max_iters):
super().__init__()
self.model_dir = model_dir
if not os.path.exists(f'checkpoints/{model_dir}'):
os.makedirs(f'checkpoints/{model_dir}')
self.logs_dir = f'checkpoints/{model_dir}/logs'
if not os.path.exists(self.logs_dir):
os.makedirs(self.logs_dir)
self.writer = SummaryWriter(self.logs_dir)
self.arcface = ArcFaceNet(50, 0.6, 'ir_se').cuda()
self.arcface.eval()
self.arcface.load_state_dict(torch.load(
'checkpoints/model_ir_se50.pth', map_location='cuda'),
strict=False)
self.mobiface = MobileFaceNet(512).cuda()
self.mobiface.eval()
self.mobiface.load_state_dict(torch.load(
'checkpoints/mobilefacenet.pth', map_location='cuda'),
strict=False)
self.generator = Generator().cuda()
self.discriminator = Discriminator().cuda()
self.adversarial_weight = 1
self.src_id_weight = 5
self.tgt_id_weight = 1
self.attributes_weight = 10
self.reconstruction_weight = 10
self.lr = lr
self.warmup = warmup
self.g_optimizer = g_optimizer(self.generator.parameters(),
lr=lr,
betas=(0, 0.999))
self.d_optimizer = d_optimizer(self.discriminator.parameters(),
lr=lr,
betas=(0, 0.999))
self.generator, self.g_optimizer = amp.initialize(self.generator,
self.g_optimizer,
opt_level="O1")
self.discriminator, self.d_optimizer = amp.initialize(
self.discriminator, self.d_optimizer, opt_level="O1")
self._iter = nn.Parameter(torch.tensor(1), requires_grad=False)
self.max_iters = max_iters
if torch.cuda.is_available():
self.cuda()
def __init__(self, config: HiDDenConfiguration, device: torch.device):
self.enc_dec = EncoderDecoder(config).to(device)
self.discr = Discriminator(config).to(device)
self.opt_enc_dec = torch.optim.Adam(self.enc_dec.parameters())
self.opt_discr = torch.optim.Adam(self.discr.parameters())
self.config = config
self.device = device
self.bce_with_logits_loss = nn.BCEWithLogitsLoss().to(device)
self.mse_loss = nn.MSELoss().to(device)
self.cover_label = 1
self.encod_label = 0
def __init__(self, configuration: HiDDenConfiguration,
device: torch.device, noiser: Noiser, tb_logger):
"""
:param configuration: Configuration for the net, such as the size of the input image, number of channels in the intermediate layers, etc.
:param device: torch.device object, CPU or GPU
:param noiser: Object representing stacked noise layers.
:param tb_logger: Optional TensorboardX logger object, if specified -- enables Tensorboard logging
"""
super(Hidden, self).__init__()
self.encoder_decoder = EncoderDecoder(configuration, noiser).to(device)
self.discriminator = Discriminator(configuration).to(device)
self.optimizer_enc_dec = torch.optim.Adam(
self.encoder_decoder.parameters())
self.optimizer_discrim = torch.optim.Adam(
self.discriminator.parameters())
if configuration.use_vgg:
self.vgg_loss = VGGLoss(3, 1, False)
self.vgg_loss.to(device)
else:
self.vgg_loss = None
self.config = configuration
self.device = device
self.bce_with_logits_loss = nn.BCEWithLogitsLoss().to(device)
self.mse_loss = nn.MSELoss().to(device)
self.ce_loss = nn.CrossEntropyLoss().to(device)
# Defined the labels used for training the discriminator/adversarial loss
self.cover_label = 1
self.encoded_label = 0
self.tb_logger = tb_logger
if tb_logger is not None:
from tensorboard_logger import TensorBoardLogger
#print(self.encoder_decoder.encoder._modules['module'].final_layer)
encoder_final = self.encoder_decoder.encoder._modules[
'module'].final_layer
encoder_final.weight.register_hook(
tb_logger.grad_hook_by_name('grads/encoder_out'))
decoder_final = self.encoder_decoder.decoder._modules[
'module'].linear
decoder_final.weight.register_hook(
tb_logger.grad_hook_by_name('grads/decoder_out'))
#print(self.discriminator._modules)
discrim_final = self.discriminator._modules['linear']
discrim_final.weight.register_hook(
tb_logger.grad_hook_by_name('grads/discrim_out'))
def __init__(self, opt):
super().__init__(opt)
self.w = opt.image_size
# make a folder for save images
self.image_dir = os.path.join(self.save_dir, 'images')
if not os.path.isdir(self.image_dir):
os.mkdir(self.image_dir)
# initialize networks
self.model_names = ['C', 'T']
self.netC = CorrespondenceNet(opt)
self.netT = TranslationNet(opt)
if opt.isTrain:
self.model_names.append('D')
self.netD = Discriminator(opt)
self.visual_names = ['b_exemplar', 'a', 'b_gen',
'b_gt'] # HPT convention
if opt.isTrain:
# assign losses
self.loss_names = [
'perc', 'domain', 'feat', 'context', 'reg', 'adv'
]
self.visual_names += ['b_warp']
self.criterionFeat = torch.nn.L1Loss()
# Both interface for VGG and perceptual loss
# call with different mode and layer params
self.criterionVGG = VGGLoss(self.device)
# Support hinge loss
self.criterionAdv = GANLoss(gan_mode=opt.gan_mode).to(self.device)
self.criterionDomain = nn.L1Loss()
self.criterionReg = torch.nn.L1Loss()
# initialize optimizers
gen_params = itertools.chain(self.netT.parameters(),
self.netC.parameters())
self.optG = torch.optim.Adam(gen_params,
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optD = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = [self.optG, self.optD]
# Finally, load checkpoints and recover schedulers
self.setup(opt)
torch.autograd.set_detect_anomaly(True)
def load_discriminator(d_params=None, ckpt_dir=None):
from model.discriminator import Discriminator
if d_params is None:
d_params = {
'labels_dim': 254,
'resolutions': [4, 8, 16, 32, 64, 128, 256, 512, 1024],
'featuremaps': [512, 512, 512, 512, 512, 256, 128, 64, 32],
}
res = d_params['resolutions'][-1]
test_images = tf.ones((1, 3, res, res), dtype=tf.float32)
# test_labels = tf.ones((1, d_params['labels_dim']), dtype=tf.float32)
# build discriminator model
discriminator = Discriminator(d_params)
_ = discriminator(test_images, training=False)
if ckpt_dir is not None:
ckpt = tf.train.Checkpoint(discriminator=discriminator)
manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=1)
ckpt.restore(manager.latest_checkpoint).expect_partial()
if manager.latest_checkpoint:
print('Discriminator restored from {}'.format(
manager.latest_checkpoint))
return discriminator
def _build_model(self):
print('start building modle.')
if self.video:
generator = Transformer(smooth_loss =self.shuffle_smooth_loss,batch_size=self.batch_size,d_model=self.features_channels, spwan=self.receptive_field, n_layers=self.attention_layer, n_head=8,big_sorting_layer = False,\
no_positional_coding =False,without_loss = 1-self.shuffle_aug, cross_entropy_loss=True)
print('created video model.')
else:
generator = get_pose_net()
print('created single-frame model.')
if self.fine_tune or self.eval:
generator = self.load_model(self.gmodel_path,generator)
self.generator = nn.DataParallel(generator).cuda()
self.optimizer = torch.optim.SGD(self.generator.parameters(), lr = self.lr, momentum = 0.9)
self.e_sche = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, milestones = self.adjust_lr_epoch, gamma = self.adjust_lr_factor)
if not self.eval:
discriminator = Discriminator()
if self.fine_tune:
discriminator = self.load_model(self.dmodel_path,discriminator)
self.discriminator = nn.DataParallel(discriminator).cuda()
self.d_opt = torch.optim.Adam(self.discriminator.parameters(),lr = self.lr)
self.d_sche = torch.optim.lr_scheduler.MultiStepLR(self.d_opt, milestones = self.adjust_lr_epoch, gamma = self.adjust_lr_factor)
if self.eval_with_single_frame_network:
spatial_feature_extractor = get_pose_net()
model_path_single_frame = config.best_model_dict['dsd_single']
spatial_feature_extractor = self.load_model(model_path_single_frame,spatial_feature_extractor)
self.spatial_feature_extractor = nn.DataParallel(spatial_feature_extractor).cuda()
self.smpl = SMPL(self.smpl_model, joint_type = 'lsp', obj_saveable = True).cuda()
print('finished build model.')
def load_gan(self):
# Get the restored generator
self.G = Generator(0, 0, restore=True)
# Use generator's latent size
self.latent_size = self.G.z_length
# Get the restored discriminator
self.D = Discriminator(0, restore=True)
def __init__(self, nc=1, nz=100, ngf=64, ndf=64, lr=0.0002, beta1=0.5, ngpu=1, autosave=None):
self.nz = nz
self.dataloader = None
self.img_list = None
self.G_losses = None
self.D_losses = None
self.device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
self.netG = Generator(nz, ngf, nc).to(self.device)
self.netG.apply(utils.weights_init)
self.netD = Discriminator(nc, ndf).to(self.device)
self.netD.apply(utils.weights_init)
self.criterion = nn.BCELoss()
self.fixed_noise = torch.randn(64, nz, 1, 1, device=self.device)
self.real_label = 1.
self.fake_label = 0.
self.optimizerD = optim.Adam(self.netD.parameters(), lr=lr, betas=(beta1, 0.999))
self.optimizerG = optim.Adam(self.netG.parameters(), lr=lr, betas=(beta1, 0.999))
self.result_path = autosave
def prepare_models(config):
discriminator = Discriminator.create_model(
config["model"]["discriminator"])
ModelInitializator._init_model(discriminator, "discriminator", config)
generator = Generator.create_model(config["model"]["generator"])
ModelInitializator._init_model(generator, "generator", config)
gan = Gan.create_model(discriminator, generator, config)
ModelInitializator._init_model(gan, "gan", config)
return generator, discriminator, gan
def __init__(self, num_D = 3, ndf = 16, n_layers = 3, downsampling_factor = 4, disc_out = 512):
super().__init__()
self.model = nn.ModuleDict()
for i in range(num_D):
self.model[f"disc_{i}"] = Discriminator(
ndf, n_layers, downsampling_factor, disc_out
)
self.downsample = nn.AvgPool1d(downsampling_factor, stride=2, padding=1, count_include_pad=False)
self.apply(weights_init)
def new(self):
print('Creating a new model...')
generator = Generator(self.resolution, self.batch_size, depth=12)
discriminator = Discriminator(self.resolution, depth=12)
lr = 5e-5
generator.model.optimizer = keras.optimizers.RMSprop(lr)
discriminator.model.optimizer = keras.optimizers.RMSprop(lr)
self.restored = False
self.initialize(generator, discriminator)
def generate_fakes(config_path, num_fake_batches, checkpoint_dir, generated_images_dir):
cfg = ModelConfig(config_path)
discriminator = Discriminator(cfg)
generator = Generator(cfg)
# Initialize models
generator(tf.ones([cfg.batch_size, cfg.latent_size]), tf.ones([cfg.batch_size, cfg.labels_size]))
discriminator(tf.ones([cfg.batch_size, cfg.resolution, cfg.resolution,
cfg.num_channels]), tf.ones([cfg.batch_size, cfg.labels_size]))
# Initialize checkpoint and manager
checkpoint = tf.train.Checkpoint(discriminator=discriminator, generator=generator)
manager = tf.train.CheckpointManager(
checkpoint, checkpoint_dir, max_to_keep=10)
if cfg.label_conditioning:
fake_labels = np.zeros([cfg.batch_size, cfg.labels_size], dtype=np.float32)
fake_labels[np.arange(cfg.batch_size), np.random.randint(cfg.labels_size, size=cfg.batch_size)] = 1.0
random_input = tf.random.normal([cfg.batch_size, cfg.latent_size])
for checkpoint_path in manager.checkpoints:
print('Restoring checkpoint from {}'.format(checkpoint_path))
checkpoint.restore(checkpoint_path).assert_consumed()
fake_image_batches = []
fake_labels_batches = []
for _ in range(num_fake_batches):
if cfg.label_conditioning:
fake_labels_batch = np.zeros([cfg.batch_size, cfg.labels_size], dtype=np.float32)
fake_labels_batch[np.arange(cfg.batch_size), np.random.randint(
cfg.labels_size, size=cfg.batch_size)] = 1.0
else:
fake_labels_batch = None
random_input = tf.random.normal([cfg.batch_size, cfg.latent_size])
fake_images_batch = generator(random_input, fake_labels_batch)
fake_image_batches.append(fake_images_batch)
fake_labels_batches.append(fake_labels_batch)
fake_images = np.concatenate(fake_image_batches, axis=0)
if cfg.label_conditioning:
fake_labels = np.concatenate(fake_labels_batches, axis=0)
else:
fake_labels = None
if not os.path.exists(generated_images_dir):
os.makedirs(generated_images_dir)
plot_batch(fake_images, fake_labels, generated_images_dir + '/checkpoint-' +
checkpoint_path.split('-')[-1], cfg.label_conditioning)
def __init__(self, HR_Shape, LR_Shape, nBaseBlocks=6, nResidualBlocks=2):
self.HR_Shape = HR_Shape
self.LR_Shape = LR_Shape
ScaleFactor = self.HR_Shape[0] // self.LR_Shape[0]
self.Generator = Generator(nBaseBlocks, nResidualBlocks,
ScaleFactor).BuildGenerator()
self.Discriminator = Discriminator(HR_Shape).BuildDiscriminator()
AdamOptimizer = Adam(lr=1E-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08)
self.Discriminator.compile(loss="binary_crossentropy",
optimizer=AdamOptimizer,
metrics=[])
self.Discriminator.trainable = False
GAN_input = Input(shape=self.LR_Shape)
x = self.Generator(GAN_input)
GAN_output = self.Discriminator(x)
self.GAN = Model(inputs=GAN_input, outputs=[x, GAN_output])
self.GAN.compile(loss=[ContentLoss, "binary_crossentropy"],
loss_weights=[0.006, 0.001],
optimizer=AdamOptimizer)
def __init__(self, hparams, device, G_AB, G_BA):
super(Model, self).__init__()
self.hparams = hparams
self.device = device
self.input_shape = hparams.input_shape
self.learning_rate = hparams.learning_rate
self.B1 = hparams.b1
self.B2 = hparams.b1
self.n_epochs = hparams.n_epochs
self.start_epoch = hparams.start_epoch
self.epoch_decay = hparams.epoch_decay
self.batch_size = hparams.batch_size
self.lambda_cycle_loss = hparams.lambda_cycle_loss
self.lambda_identity_loss = hparams.lambda_identity_loss
self.G_AB = G_AB
self.G_BA = G_BA
self.D_A = Discriminator(self.input_shape)
self.D_B = Discriminator(self.input_shape)
# Adversarial ground truths
self.valid = torch.ones(
(self.batch_size, *self.D_A.output_shape)).to(device)
self.fake = torch.zeros(
(self.batch_size, *self.D_A.output_shape)).to(device)
# Losses
self.criterion_GAN = torch.nn.MSELoss()
self.criterion_cycle = torch.nn.L1Loss()
self.criterion_identity = torch.nn.L1Loss()
self.fake_A = None
self.fake_B = None
self.recov_A = None
self.recov_B = None
def build_base_component(self):
# hyper params
self.lr = tf.placeholder(tf.float32, shape=(),
name='lr') # learning rate
self.dropout = tf.placeholder(tf.float32, shape=(), name='dropout')
# input of the graph
self.img = tf.placeholder(tf.uint8,
shape=(None, None, None, 1),
name='img') # (N, H, W, C)，这里C=1，因为是灰度图
self.formula = tf.placeholder(tf.int32,
shape=(None, None),
name='formula') # (N, formula_tokens)
self.formula_length = tf.placeholder(tf.int32,
shape=(None, ),
name='formula_length') # (N, 1)
# self.pred_train, self.pred_test
# tensorflow 只有静态计算图，只好同时把 train 和 test 部分的计算图都建了
self.generator = Generator(self._config, self._vocab.n_tok,
self._vocab.id_end)
train, test = self.generator(self.img, self.formula, self.dropout)
self.pred_train = train
self.pred_test = test
self.discriminator = Discriminator(self._config, self._vocab.n_tok)
self.D_loss = self.discriminator(self.pred_test.ids, self.formula,
self.dropout)
self.D_optimizer = tf.train.AdamOptimizer().minimize(self.D_loss)
# self.loss 生成器第一阶段的 loss
losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=self.pred_train, labels=self.formula)
mask = tf.sequence_mask(self.formula_length)
losses = tf.boolean_mask(losses, mask)
self.loss = tf.reduce_mean(losses)
# to compute perplexity for test
self.ce_words = tf.reduce_sum(losses) # sum of CE for each word
self.n_words = tf.reduce_sum(self.formula_length) # number of words
# tensorboard
tf.summary.image("img", self.img)
tf.summary.scalar("learning_rate", self.lr)
tf.summary.scalar("dropout", self.dropout)
tf.summary.scalar("G_loss", self.loss)
tf.summary.scalar("D_loss", self.D_loss)
tf.summary.scalar("sum_of_CE_for_each_word", self.ce_words)
tf.summary.scalar("number_of_words", self.n_words)
class Network(nn.Module):
def __init__(self):
super(Network, self).__init__()
self.feature_extractor = resnet50() # Already pretrained
# self.feature_extractor = resnet50(pretrained_path=None)
self.selector = Selector()
self.dis = Discriminator()
self.optmzr_select = Adam(self.selector.parameters(), lr=1e-3)
self.optmzr_dis = Adam(self.dis.parameters(), lr=1e-3)
def forward(self, anchor: Variable, real_data: Variable, fake_data: Variable):
assert len(anchor.size()) == 4 and len(anchor.size()) == 4
fea_anchor = self.feature_extractor(anchor)
fea_real = self.feature_extractor(real_data)
fea_fake = self.feature_extractor(fake_data)
# not train_feature:
fea_anchor = fea_anchor.detach()
fea_real = fea_real.detach()
fea_fake = fea_fake.detach()
score_real = self.dis(fea_anchor, fea_real)
score_fake = self.dis(fea_anchor, fea_fake)
return score_real, score_fake
def bp_dis(self, score_real, score_fake):
real_label = Variable(torch.normal(torch.ones(score_real.size()), torch.zeros(score_real.size()) + 0.05)).cuda()
fake_label = Variable(
torch.normal(torch.zeros(score_real.size()), torch.zeros(score_real.size()) + 0.05)).cuda()
loss = torch.mean(F.binary_cross_entropy(score_real, real_label, size_average=False) + \
F.binary_cross_entropy(score_fake, fake_label, size_average=False))
# loss = -(torch.mean(torch.log(score_real + 1e-6)) - torch.mean(torch.log(.5 + score_fake / 2 + 1e-6)))
self.optmzr_dis.zero_grad()
loss.backward()
return self.optmzr_dis.step()
def bp_select(self, score_fake: Variable, fake_prob):
# torch.mean(torch.log(prob) * torch.log(1 - score_fake), 0)
n_sample = score_fake.size()[0]
self.optmzr_dis.zero_grad()
re = (score_fake.data - .5) * 2
torch.log(fake_prob).backward(re / n_sample)
def init_models(self):
generator = Generator(self.config).to(self.config.device)
generator.apply(weights_init)
if self.config.generator_path is not None:
generator.load_state_dict(torch.load(self.config.generator_path))
# print(generator)
discriminator = Discriminator(self.config).to(self.config.device)
discriminator.apply(weights_init)
if self.config.discriminator_path is not None:
discriminator.load_state_dict(
torch.load(self.config.discriminator_path))
# print(discriminator)
return discriminator, generator
def main():
# set torch and numpy seed for reproducibility
torch.manual_seed(27)
np.random.seed(27)
# tensorboard writer
writer = SummaryWriter(settings.TENSORBOARD_DIR)
# makedir snapshot
makedir(settings.CHECKPOINT_DIR)
# enable cudnn
torch.backends.cudnn.enabled = True
# create segmentor network
model_G = Segmentor(pretrained=settings.PRETRAINED,
num_classes=settings.NUM_CLASSES,
modality=settings.MODALITY)
model_G.train()
model_G.cuda()
torch.backends.cudnn.benchmark = True
# create discriminator network
model_D = Discriminator(settings.NUM_CLASSES)
model_D.train()
model_D.cuda()
# dataset and dataloader
dataset = TrainDataset()
dataloader = data.DataLoader(dataset,
batch_size=settings.BATCH_SIZE,
shuffle=True,
num_workers=settings.NUM_WORKERS,
pin_memory=True,
drop_last=True)
test_dataset = TestDataset(data_root=settings.DATA_ROOT_VAL,
data_list=settings.DATA_LIST_VAL)
test_dataloader = data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=settings.NUM_WORKERS,
pin_memory=True)
# optimizer for generator network (segmentor)
optim_G = optim.SGD(model_G.optim_parameters(settings.LR),
lr=settings.LR,
momentum=settings.LR_MOMENTUM,
weight_decay=settings.WEIGHT_DECAY)
# lr scheduler for optimi_G
lr_lambda_G = lambda epoch: (1 - epoch / settings.EPOCHS
)**settings.LR_POLY_POWER
lr_scheduler_G = optim.lr_scheduler.LambdaLR(optim_G,
lr_lambda=lr_lambda_G)
# optimizer for discriminator network
optim_D = optim.Adam(model_D.parameters(), settings.LR_D)
# lr scheduler for optimi_D
lr_lambda_D = lambda epoch: (1 - epoch / settings.EPOCHS
)**settings.LR_POLY_POWER
lr_scheduler_D = optim.lr_scheduler.LambdaLR(optim_D,
lr_lambda=lr_lambda_D)
# losses
ce_loss = CrossEntropyLoss2d(
ignore_index=settings.IGNORE_LABEL) # to use for segmentor
bce_loss = BCEWithLogitsLoss2d() # to use for discriminator
# upsampling for the network output
upsample = nn.Upsample(size=(settings.CROP_SIZE, settings.CROP_SIZE),
mode='bilinear',
align_corners=True)
# # labels for adversarial training
# pred_label = 0
# gt_label = 1
# load the model to resume training
last_epoch = -1
if settings.RESUME_TRAIN:
checkpoint = torch.load(settings.LAST_CHECKPOINT)
model_G.load_state_dict(checkpoint['model_G_state_dict'])
model_G.train()
model_G.cuda()
model_D.load_state_dict(checkpoint['model_D_state_dict'])
model_D.train()
model_D.cuda()
optim_G.load_state_dict(checkpoint['optim_G_state_dict'])
optim_D.load_state_dict(checkpoint['optim_D_state_dict'])
lr_scheduler_G.load_state_dict(checkpoint['lr_scheduler_G_state_dict'])
lr_scheduler_D.load_state_dict(checkpoint['lr_scheduler_D_state_dict'])
last_epoch = checkpoint['epoch']
# purge the logs after the last_epoch
writer = SummaryWriter(settings.TENSORBOARD_DIR,
purge_step=(last_epoch + 1) * len(dataloader))
for epoch in range(last_epoch + 1, settings.EPOCHS + 1):
train_one_epoch(model_G,
model_D,
optim_G,
optim_D,
dataloader,
test_dataloader,
epoch,
upsample,
ce_loss,
bce_loss,
writer,
print_freq=5,
eval_freq=settings.EVAL_FREQ)
if epoch % settings.CHECKPOINT_FREQ == 0 and epoch != 0:
save_checkpoint(epoch, model_G, model_D, optim_G, optim_D,
lr_scheduler_G, lr_scheduler_D)
# save the final model
if epoch >= settings.EPOCHS:
print('saving the final model')
save_checkpoint(epoch, model_G, model_D, optim_G, optim_D,
lr_scheduler_G, lr_scheduler_D)
writer.close()
lr_scheduler_G.step()
lr_scheduler_D.step()
def __init__(self, config, args):
self.args = args
self.config = config
self.visdom = args.visdom
if args.visdom:
self.vis = visdom.Visdom(env=os.getcwd().split('/')[-1], port=8888)
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
config)
self.target_train_loader, self.target_val_loader, self.target_test_loader, _ = make_target_data_loader(
config)
# Define network
self.model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
self.D = Discriminator(num_classes=self.nclass, ndf=16)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': config.lr
}, {
'params': self.model.get_10x_lr_params(),
'lr': config.lr * config.lr_ratio
}]
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.D_optimizer = torch.optim.Adam(self.D.parameters(),
lr=config.lr,
betas=(0.9, 0.99))
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(
weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.entropy_mini_loss = MinimizeEntropyLoss()
self.bottleneck_loss = BottleneckLoss()
self.instance_loss = InstanceLoss()
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
self.summary = TensorboardSummary('./train_log')
# labels for adversarial training
self.source_label = 0
self.target_label = 1
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.D = torch.nn.DataParallel(self.D)
patch_replication_callback(self.D)
self.D = self.D.cuda()
self.best_pred_source = 0.0
self.best_pred_target = 0.0
# Resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
class Trainer(object):
def __init__(self, config, args):
self.args = args
self.config = config
self.visdom = args.visdom
if args.visdom:
self.vis = visdom.Visdom(env=os.getcwd().split('/')[-1], port=8888)
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
config)
self.target_train_loader, self.target_val_loader, self.target_test_loader, _ = make_target_data_loader(
config)
# Define network
self.model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
self.D = Discriminator(num_classes=self.nclass, ndf=16)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': config.lr
}, {
'params': self.model.get_10x_lr_params(),
'lr': config.lr * config.lr_ratio
}]
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.D_optimizer = torch.optim.Adam(self.D.parameters(),
lr=config.lr,
betas=(0.9, 0.99))
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(
weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.entropy_mini_loss = MinimizeEntropyLoss()
self.bottleneck_loss = BottleneckLoss()
self.instance_loss = InstanceLoss()
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
self.summary = TensorboardSummary('./train_log')
# labels for adversarial training
self.source_label = 0
self.target_label = 1
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.D = torch.nn.DataParallel(self.D)
patch_replication_callback(self.D)
self.D = self.D.cuda()
self.best_pred_source = 0.0
self.best_pred_target = 0.0
# Resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
def training(self, epoch):
train_loss, seg_loss_sum, bn_loss_sum, entropy_loss_sum, adv_loss_sum, d_loss_sum, ins_loss_sum = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
self.model.train()
if config.freeze_bn:
self.model.module.freeze_bn()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
target_train_iterator = iter(self.target_train_loader)
for i, sample in enumerate(tbar):
itr = epoch * len(self.train_loader) + i
#if self.visdom:
# self.vis.line(X=torch.tensor([itr]), Y=torch.tensor([self.optimizer.param_groups[0]['lr']]),
# win='lr', opts=dict(title='lr', xlabel='iter', ylabel='lr'),
# update='append' if itr>0 else None)
self.summary.writer.add_scalar(
'Train/lr', self.optimizer.param_groups[0]['lr'], itr)
A_image, A_target = sample['image'], sample['label']
# Get one batch from target domain
try:
target_sample = next(target_train_iterator)
except StopIteration:
target_train_iterator = iter(self.target_train_loader)
target_sample = next(target_train_iterator)
B_image, B_target, B_image_pair = target_sample[
'image'], target_sample['label'], target_sample['image_pair']
if self.args.cuda:
A_image, A_target = A_image.cuda(), A_target.cuda()
B_image, B_target, B_image_pair = B_image.cuda(
), B_target.cuda(), B_image_pair.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred_source,
self.best_pred_target, self.config.lr_ratio)
self.scheduler(self.D_optimizer, i, epoch, self.best_pred_source,
self.best_pred_target, self.config.lr_ratio)
A_output, A_feat, A_low_feat = self.model(A_image)
B_output, B_feat, B_low_feat = self.model(B_image)
#B_output_pair, B_feat_pair, B_low_feat_pair = self.model(B_image_pair)
#B_output_pair, B_feat_pair, B_low_feat_pair = flip(B_output_pair, dim=-1), flip(B_feat_pair, dim=-1), flip(B_low_feat_pair, dim=-1)
self.optimizer.zero_grad()
self.D_optimizer.zero_grad()
# Train seg network
for param in self.D.parameters():
param.requires_grad = False
# Supervised loss
seg_loss = self.criterion(A_output, A_target)
main_loss = seg_loss
# Unsupervised loss
#ins_loss = 0.01 * self.instance_loss(B_output, B_output_pair)
#main_loss += ins_loss
# Train adversarial loss
D_out = self.D(prob_2_entropy(F.softmax(B_output)))
adv_loss = bce_loss(D_out, self.source_label)
main_loss += self.config.lambda_adv * adv_loss
main_loss.backward()
# Train discriminator
for param in self.D.parameters():
param.requires_grad = True
A_output_detach = A_output.detach()
B_output_detach = B_output.detach()
# source
D_source = self.D(prob_2_entropy(F.softmax(A_output_detach)))
source_loss = bce_loss(D_source, self.source_label)
source_loss = source_loss / 2
# target
D_target = self.D(prob_2_entropy(F.softmax(B_output_detach)))
target_loss = bce_loss(D_target, self.target_label)
target_loss = target_loss / 2
d_loss = source_loss + target_loss
d_loss.backward()
self.optimizer.step()
self.D_optimizer.step()
seg_loss_sum += seg_loss.item()
#ins_loss_sum += ins_loss.item()
adv_loss_sum += self.config.lambda_adv * adv_loss.item()
d_loss_sum += d_loss.item()
#train_loss += seg_loss.item() + self.config.lambda_adv * adv_loss.item()
train_loss += seg_loss.item()
self.summary.writer.add_scalar('Train/SegLoss', seg_loss.item(),
itr)
#self.summary.writer.add_scalar('Train/InsLoss', ins_loss.item(), itr)
self.summary.writer.add_scalar('Train/AdvLoss', adv_loss.item(),
itr)
self.summary.writer.add_scalar('Train/DiscriminatorLoss',
d_loss.item(), itr)
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
# Show the results of the last iteration
#if i == len(self.train_loader)-1:
print("Add Train images at epoch" + str(epoch))
self.summary.visualize_image('Train-Source', self.config.dataset,
A_image, A_target, A_output, epoch, 5)
self.summary.visualize_image('Train-Target', self.config.target,
B_image, B_target, B_output, epoch, 5)
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + A_image.data.shape[0]))
print('Loss: %.3f' % train_loss)
#print('Seg Loss: %.3f' % seg_loss_sum)
#print('Ins Loss: %.3f' % ins_loss_sum)
#print('BN Loss: %.3f' % bn_loss_sum)
#print('Adv Loss: %.3f' % adv_loss_sum)
#print('Discriminator Loss: %.3f' % d_loss_sum)
#if self.visdom:
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([seg_loss_sum]), win='train_loss', name='Seg_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([ins_loss_sum]), win='train_loss', name='Ins_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([bn_loss_sum]), win='train_loss', name='BN_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([adv_loss_sum]), win='train_loss', name='Adv_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([d_loss_sum]), win='train_loss', name='Dis_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
def validation(self, epoch):
def get_metrics(tbar, if_source=False):
self.evaluator.reset()
test_loss = 0.0
#feat_mean, low_feat_mean, feat_var, low_feat_var = 0, 0, 0, 0
#adv_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output, low_feat, feat = self.model(image)
#low_feat = low_feat.cpu().numpy()
#feat = feat.cpu().numpy()
#if isinstance(feat, np.ndarray):
# feat_mean += feat.mean(axis=0).mean(axis=1).mean(axis=1)
# low_feat_mean += low_feat.mean(axis=0).mean(axis=1).mean(axis=1)
# feat_var += feat.var(axis=0).var(axis=1).var(axis=1)
# low_feat_var += low_feat.var(axis=0).var(axis=1).var(axis=1)
#else:
# feat_mean = feat.mean(axis=0).mean(axis=1).mean(axis=1)
# low_feat_mean = low_feat.mean(axis=0).mean(axis=1).mean(axis=1)
# feat_var = feat.var(axis=0).var(axis=1).var(axis=1)
# low_feat_var = low_feat.var(axis=0).var(axis=1).var(axis=1)
#d_output = self.D(prob_2_entropy(F.softmax(output)))
#adv_loss += bce_loss(d_output, self.source_label).item()
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target_ = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target_, pred)
if if_source:
print("Add Validation-Source images at epoch" + str(epoch))
self.summary.visualize_image('Val-Source', self.config.dataset,
image, target, output, epoch, 5)
else:
print("Add Validation-Target images at epoch" + str(epoch))
self.summary.visualize_image('Val-Target', self.config.target,
image, target, output, epoch, 5)
#feat_mean /= (i+1)
#low_feat_mean /= (i+1)
#feat_var /= (i+1)
#low_feat_var /= (i+1)
#adv_loss /= (i+1)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
if if_source:
print('Validation on source:')
else:
print('Validation on target:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + image.data.shape[0]))
print("Acc:{}, IoU:{}, mIoU:{}".format(Acc, IoU, mIoU))
print('Loss: %.3f' % test_loss)
if if_source:
names = ['source', 'source_acc', 'source_IoU', 'source_mIoU']
self.summary.writer.add_scalar('Val/SourceAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/SourceIoU', IoU, epoch)
else:
names = ['target', 'target_acc', 'target_IoU', 'target_mIoU']
self.summary.writer.add_scalar('Val/TargetAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/TargetIoU', IoU, epoch)
# Draw Visdom
#if if_source:
# names = ['source', 'source_acc', 'source_IoU', 'source_mIoU']
#else:
# names = ['target', 'target_acc', 'target_IoU', 'target_mIoU']
#if self.visdom:
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([test_loss]), win='val_loss', name=names[0],
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([adv_loss]), win='val_loss', name='adv_loss',
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([Acc]), win='metrics', name=names[1],
# opts=dict(title='metrics', xlabel='epoch', ylabel='performance'),
# update='append' if epoch > 0 else None)
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([IoU]), win='metrics', name=names[2],
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([mIoU]), win='metrics', name=names[3],
# update='append')
return Acc, IoU, mIoU
self.model.eval()
tbar_source = tqdm(self.val_loader, desc='\r')
tbar_target = tqdm(self.target_val_loader, desc='\r')
s_acc, s_iou, s_miou = get_metrics(tbar_source, True)
t_acc, t_iou, t_miou = get_metrics(tbar_target, False)
new_pred_source = s_iou
new_pred_target = t_iou
if new_pred_source > self.best_pred_source or new_pred_target > self.best_pred_target:
is_best = True
self.best_pred_source = max(new_pred_source, self.best_pred_source)
self.best_pred_target = max(new_pred_target, self.best_pred_target)
print('Saving state, epoch:', epoch)
torch.save(
self.model.module.state_dict(),
self.args.save_folder + 'models/' + 'epoch' + str(epoch) + '.pth')
loss_file = {
's_Acc': s_acc,
's_IoU': s_iou,
's_mIoU': s_miou,
't_Acc': t_acc,
't_IoU': t_iou,
't_mIoU': t_miou
}
with open(
os.path.join(self.args.save_folder, 'eval',
'epoch' + str(epoch) + '.json'), 'w') as f:
json.dump(loss_file, f)
def train_model(config_path, data_path):
cfg = ModelConfig(config_path)
discriminator = Discriminator(cfg)
generator = Generator(cfg)
gen_optimizer = tf.keras.optimizers.Adam(
learning_rate=cfg.generator_base_learning_rate,
beta_1=cfg.generator_beta_1,
beta_2=cfg.generator_beta_2,
epsilon=1e-8)
disc_optimizer = tf.keras.optimizers.Adam(
learning_rate=cfg.discriminator_base_learning_rate,
beta_1=cfg.discriminator_beta_1,
beta_2=cfg.discriminator_beta_2,
epsilon=1e-8)
dataset = read_dataset(data_path, cfg)
# Initialize checkpoint and checkpoint manager
ckpt_models = tf.train.Checkpoint(generator=generator,
discriminator=discriminator)
ckpt_optimizers = tf.train.Checkpoint(gen_optimizer=gen_optimizer,
disc_optimizer=disc_optimizer)
ckpt_manager_models = tf.train.CheckpointManager(
checkpoint=ckpt_models,
directory=cfg.checkpoint_path + '/models/',
max_to_keep=cfg.max_checkpoints_to_keep)
ckpt_manager_optimizers = tf.train.CheckpointManager(
checkpoint=ckpt_optimizers,
directory=cfg.checkpoint_path + '/optimizers/',
max_to_keep=cfg.max_checkpoints_to_keep)
# Initialize log writer
train_summary_writer = tf.summary.create_file_writer(cfg.log_dir)
# Initialize metrics
gen_loss = tf.keras.metrics.Metric
start_time = time.time()
num_images_before = 0
num_minibatch = 0
for example in dataset:
disc_regularization = (num_minibatch % cfg.disc_reg_intervall == 0)
gen_loss, disc_loss = train_step(
real_images=example['data'],
real_labels=example['label'],
generator=generator,
discriminator=discriminator,
gen_optimizer=gen_optimizer,
disc_optimizer=disc_optimizer,
cfg=cfg,
disc_regularization=disc_regularization)
num_minibatch = gen_optimizer.iterations.numpy()
num_images = num_minibatch * cfg.batch_size
# Print Metrics
if (num_images %
(cfg.print_metrics_intervall_kimg * 1000)) < cfg.batch_size:
images_per_second = (num_images - num_images_before) / (
time.time() - start_time)
print('minibatch {} images {} gen loss {:.4f} disc loss {:.4f}'
' images per second {:.2f}'.format(num_minibatch, num_images,
gen_loss, disc_loss,
images_per_second))
num_images_before = num_images
start_time = time.time()
# Save checkpoint
if (num_images %
(cfg.checkpoint_intervall_kimg * 1000)) < cfg.batch_size:
save_checkpoint(ckpt_manager_models, num_images)
save_checkpoint(ckpt_manager_optimizers, num_images)
# Log metrics
if (num_images %
(cfg.log_metrics_intervall_kimg * 1000)) < cfg.batch_size:
with train_summary_writer.as_default():
tf.summary.scalar('gen_loss', gen_loss, step=num_images)
tf.summary.scalar('disc_loss', disc_loss, step=num_images)
if (num_images % (cfg.max_num_images_kimg * 1000)) < cfg.batch_size:
# Save final state if not already done
if not (num_images %
cfg.checkpoint_intervall_kimg) < cfg.batch_size:
save_checkpoint(ckpt_manager_models, int(num_images / 1000))
save_checkpoint(ckpt_manager_optimizers,
int(num_images / 1000))
break
def __init__(self,
width,
height,
xchan,
ychan,
lambda_=10.,
pool_size=50,
lr=0.0002,
beta1=0.5):
"""
width: image width in pixel.
height: image height in pixel.
ichan: number of channels used by input images.
ochan: number of channels used by output images.
lambda_: Cycle-Consistency weighting.
pool_size: Image pool size.
lr: learning rate for ADAM optimizer.
beta1: beta1 parameter for ADAM optimizer.
"""
self._dx_pool = _ImagePool(pool_size)
self._dy_pool = _ImagePool(pool_size)
self._xs = tf.placeholder(tf.float32, [None, width, height, xchan])
self._ys = tf.placeholder(tf.float32, [None, width, height, ychan])
self._d_xs = tf.placeholder(tf.float32, [None, width, height, xchan])
self._d_ys = tf.placeholder(tf.float32, [None, width, height, ychan])
self._fake_d_xs = tf.placeholder(tf.float32,
[None, width, height, xchan])
self._fake_d_ys = tf.placeholder(tf.float32,
[None, width, height, ychan])
self._gx = Generator('Gx', self._ys, xchan)
self._gy = Generator('Gy', self._xs, ychan)
self._gx_from_gy = Generator('Gx',
self._gy['l15']['fmap'],
xchan,
reuse=True)
self._gy_from_gx = Generator('Gy',
self._gx['l15']['fmap'],
ychan,
reuse=True)
self._real_dx = Discriminator('Dx', self._d_xs)
self._fake_dx = Discriminator('Dx', self._xs, reuse=True)
self._fake_dx_g = Discriminator('Dx',
self._gx['l15']['fmap'],
reuse=True)
self._real_dy = Discriminator('Dy', self._d_ys)
self._fake_dy = Discriminator('Dy', self._ys, reuse=True)
self._fake_dy_g = Discriminator('Dy',
self._gy['l15']['fmap'],
reuse=True)
# Forward and backward Cycle-Consistency with LSGAN-kind losses
cycle_loss = lambda_ * (
tf.reduce_mean(tf.abs(
(self._gx_from_gy['l15']['fmap'] - self._xs))) +
tf.reduce_mean(tf.abs(
(self._gy_from_gx['l15']['fmap'] - self._ys))))
self._gx_loss = 0.5 * tf.reduce_mean(
tf.square(self._fake_dx_g['l5']['fmap'] - 1.)) + cycle_loss
self._gy_loss = 0.5 * tf.reduce_mean(
tf.square(self._fake_dy_g['l5']['fmap'] - 1.)) + cycle_loss
self._dx_loss = 0.5 * tf.reduce_mean(
tf.square(self._real_dx['l5']['fmap'] - 1.)
) + 0.5 * tf.reduce_mean(tf.square(self._fake_dx['l5']['fmap']))
self._dy_loss = 0.5 * tf.reduce_mean(
tf.square(self._real_dy['l5']['fmap'] - 1.)
) + 0.5 * tf.reduce_mean(tf.square(self._fake_dy['l5']['fmap']))
self._gx_train_step = tf.train.AdamOptimizer(lr, beta1=beta1).minimize(
self._gx_loss,
var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='Gx'))
self._gy_train_step = tf.train.AdamOptimizer(lr, beta1=beta1).minimize(
self._gy_loss,
var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='Gy'))
self._dx_train_step = tf.train.AdamOptimizer(lr, beta1=beta1).minimize(
self._dx_loss,
var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='Dx'))
self._dy_train_step = tf.train.AdamOptimizer(lr, beta1=beta1).minimize(
self._dy_loss,
var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='Dy'))
def main():
# define actor/critic/discriminator net and optimizer
policy = Policy(discrete_action_sections, discrete_state_sections)
value = Value()
discriminator = Discriminator()
optimizer_policy = torch.optim.Adam(policy.parameters(), lr=args.policy_lr)
optimizer_value = torch.optim.Adam(value.parameters(), lr=args.value_lr)
optimizer_discriminator = torch.optim.Adam(discriminator.parameters(),
lr=args.discrim_lr)
discriminator_criterion = nn.BCELoss()
writer = SummaryWriter()
# load expert data
dataset = ExpertDataSet(args.expert_activities_data_path,
args.expert_cost_data_path)
data_loader = data.DataLoader(dataset=dataset,
batch_size=args.expert_batch_size,
shuffle=False,
num_workers=1)
# load models
# discriminator.load_state_dict(torch.load('./model_pkl/Discriminator_model_3.pkl'))
# policy.transition_net.load_state_dict(torch.load('./model_pkl/Transition_model_3.pkl'))
# policy.policy_net.load_state_dict(torch.load('./model_pkl/Policy_model_3.pkl'))
# value.load_state_dict(torch.load('./model_pkl/Value_model_3.pkl'))
print('############# start training ##############')
# update discriminator
num = 0
for ep in tqdm(range(args.training_epochs)):
# collect data from environment for ppo update
start_time = time.time()
memory = policy.collect_samples(args.ppo_buffer_size, size=10000)
# print('sample_data_time:{}'.format(time.time()-start_time))
batch = memory.sample()
continuous_state = torch.stack(
batch.continuous_state).squeeze(1).detach()
discrete_action = torch.stack(
batch.discrete_action).squeeze(1).detach()
continuous_action = torch.stack(
batch.continuous_action).squeeze(1).detach()
next_discrete_state = torch.stack(
batch.next_discrete_state).squeeze(1).detach()
next_continuous_state = torch.stack(
batch.next_continuous_state).squeeze(1).detach()
old_log_prob = torch.stack(batch.old_log_prob).detach()
mask = torch.stack(batch.mask).squeeze(1).detach()
discrete_state = torch.stack(batch.discrete_state).squeeze(1).detach()
d_loss = torch.empty(0, device=device)
p_loss = torch.empty(0, device=device)
v_loss = torch.empty(0, device=device)
gen_r = torch.empty(0, device=device)
expert_r = torch.empty(0, device=device)
for _ in range(1):
for expert_state_batch, expert_action_batch in data_loader:
gen_state = torch.cat((discrete_state, continuous_state),
dim=-1)
gen_action = torch.cat((discrete_action, continuous_action),
dim=-1)
gen_r = discriminator(gen_state, gen_action)
expert_r = discriminator(expert_state_batch,
expert_action_batch)
optimizer_discriminator.zero_grad()
d_loss = discriminator_criterion(gen_r,
torch.zeros(gen_r.shape, device=device)) + \
discriminator_criterion(expert_r,
torch.ones(expert_r.shape, device=device))
total_d_loss = d_loss - 10 * torch.var(gen_r.to(device))
d_loss.backward()
# total_d_loss.backward()
optimizer_discriminator.step()
writer.add_scalar('d_loss', d_loss, ep)
# writer.add_scalar('total_d_loss', total_d_loss, ep)
writer.add_scalar('expert_r', expert_r.mean(), ep)
# update PPO
gen_r = discriminator(
torch.cat((discrete_state, continuous_state), dim=-1),
torch.cat((discrete_action, continuous_action), dim=-1))
optimize_iter_num = int(
math.ceil(discrete_state.shape[0] / args.ppo_mini_batch_size))
for ppo_ep in range(args.ppo_optim_epoch):
for i in range(optimize_iter_num):
num += 1
index = slice(
i * args.ppo_mini_batch_size,
min((i + 1) * args.ppo_mini_batch_size,
discrete_state.shape[0]))
discrete_state_batch, continuous_state_batch, discrete_action_batch, continuous_action_batch, \
old_log_prob_batch, mask_batch, next_discrete_state_batch, next_continuous_state_batch, gen_r_batch = \
discrete_state[index], continuous_state[index], discrete_action[index], continuous_action[index], \
old_log_prob[index], mask[index], next_discrete_state[index], next_continuous_state[index], gen_r[
index]
v_loss, p_loss = ppo_step(
policy, value, optimizer_policy, optimizer_value,
discrete_state_batch, continuous_state_batch,
discrete_action_batch, continuous_action_batch,
next_discrete_state_batch, next_continuous_state_batch,
gen_r_batch, old_log_prob_batch, mask_batch,
args.ppo_clip_epsilon)
writer.add_scalar('p_loss', p_loss, num)
writer.add_scalar('v_loss', v_loss, num)
writer.add_scalar('gen_r', gen_r.mean(), num)
print('#' * 5 + 'training episode:{}'.format(ep) + '#' * 5)
print('d_loss', d_loss.item())
# print('p_loss', p_loss.item())
# print('v_loss', v_loss.item())
print('gen_r:', gen_r.mean().item())
print('expert_r:', expert_r.mean().item())
memory.clear_memory()
# save models
torch.save(discriminator.state_dict(),
'./model_pkl/Discriminator_model_4.pkl')
torch.save(policy.transition_net.state_dict(),
'./model_pkl/Transition_model_4.pkl')
torch.save(policy.policy_net.state_dict(),
'./model_pkl/Policy_model_4.pkl')
torch.save(value.state_dict(), './model_pkl/Value_model_4.pkl')
def main():
# dataset preparation
source_data, target_data, val_data = create_dataset(mode='G2C')
source_dataloader = Data.DataLoader(source_data,
batch_size=parser.batch_size,
shuffle=True,
num_workers=parser.num_workers,
pin_memory=True)
target_dataloader = Data.DataLoader(target_data,
batch_size=parser.batch_size,
shuffle=True,
num_workers=parser.num_workers,
pin_memory=True)
val_dataloader = Data.DataLoader(val_data,
batch_size=parser.batch_size,
shuffle=False,
num_workers=parser.num_workers,
pin_memory=True)
source_dataloader_iter = enumerate(source_dataloader)
target_dataloader_iter = enumerate(target_dataloader)
save_dir = parser.ckpt_dir
# create model and optimizer
model = create_model(num_classes=parser.num_classes, name='DeepLab')
D1 = Discriminator(num_classes=parser.num_classes)
D2 = Discriminator(num_classes=parser.num_classes)
optimizer_G = create_optimizer(model.get_optim_params(parser),
lr=parser.learning_rate,
momentum=parser.momentum,
weight_decay=parser.weight_decay,
name="SGD")
optimizer_D1 = create_optimizer(D1.parameters(),
lr=LEARNING_RATE_D,
name="Adam",
betas=BETAS)
optimizer_D2 = create_optimizer(D2.parameters(),
lr=LEARNING_RATE_D,
name="Adam",
betas=BETAS)
optimizer_G.zero_grad()
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
start_iter = 1
last_mIoU = 0
if parser.restore:
print("loading checkpoint...")
checkpoint = torch.load(save_dir)
start_iter = checkpoint['iter']
model.load_state_dict(checkpoint['model'])
optimizer_G.load_state_dict(checkpoint['optimizer']['G'])
optimizer_D1.load_state_dict(checkpoint['optimizer']['D1'])
optimizer_D2.load_state_dict(checkpoint['optimizer']['D2'])
last_mIoU = checkpoint['best_mIoU']
print("start training...")
print("pytorch version: " + TORCH_VERSION + ", cuda version: " +
TORCH_CUDA_VERSION + ", cudnn version: " + CUDNN_VERSION)
print("available graphical device: " + DEVICE_NAME)
os.system("nvidia-smi")
discriminator = {'D1': D1, 'D2': D2}
optimizer = {'G': optimizer_G, 'D1': optimizer_D1, 'D2': optimizer_D2}
best_mIoU, best_iter = train(model, discriminator, optimizer,
source_dataloader_iter,
target_dataloader_iter, val_dataloader,
start_iter, last_mIoU)
print("finished training, the best mIoU is: " + str(best_mIoU) +
" in iteration " + str(best_iter))
