diff --git a/models/roma_model.py b/models/roma_model.py index a3a7af7..12536b9 100644 --- a/models/roma_model.py +++ b/models/roma_model.py @@ -68,9 +68,6 @@ class ROMAModel(BaseModel): # From UNSB self.netE = networks.define_D(opt.output_nc*4, opt.ndf, opt.netD, opt.n_layers_D, opt.normD, opt.init_type, opt.init_gain, opt.no_antialias, self.gpu_ids, opt) - - # Deine another generator - self.netG_2 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.normG, not opt.no_dropout, opt.init_type, opt.init_gain, opt.no_antialias, opt.no_antialias_up, self.gpu_ids, opt) self.norm = F.softmax @@ -186,8 +183,6 @@ class ROMAModel(BaseModel): # 保存去噪后的中间结果 (real_A_noisy 等)，供下一步做拼接 self.real_A_noisy = Xt.detach() self.real_A_noisy2 = Xt2.detach() - # 保存noisy_map - self.noisy_map = self.real_A_noisy - self.real_A # ============ 第三步：拼接输入并执行网络推理 ============= bs = self.real_A0.size(0) @@ -206,7 +201,23 @@ class ROMAModel(BaseModel): self.fake_B0 = self.netG(self.real_A0) self.fake_B1 = self.netG(self.real_A1) - + + if self.opt.phase == 'train': + # 生成图像的梯度 + fake_gradient = torch.autograd.grad(self.fake_B0.sum(), self.fake_B0, create_graph=True)[0] + # 梯度图 + self.weight_fake = self.cao.generate_weight_map(fake_gradient) + + # 生成图像的CTN光流图 + self.f_content = self.ctn(self.weight_fake) + + # 变换后的图片 + self.warped_real_A_noisy2 = warp(self.real_A_noisy, self.f_content) + self.warped_fake_B0 = warp(self.fake_B0,self.f_content) + + # 经过第二次生成器 + self.warped_fake_B0_2 = self.netG(self.warped_real_A_noisy2, self.time, z_in) + if self.opt.isTrain: real_A0 = self.real_A0 real_A1 = self.real_A1 @@ -214,98 +225,41 @@ class ROMAModel(BaseModel): real_B1 = self.real_B1 fake_B0 = self.fake_B0 fake_B1 = self.fake_B1 + warped_fake_B0_2=self.warped_fake_B0_2 + warped_fake_B0=self.warped_fake_B0 + self.real_A0_resize = self.resize(real_A0) self.real_A1_resize = self.resize(real_A1) real_B0 = self.resize(real_B0) real_B1 = self.resize(real_B1) self.fake_B0_resize = self.resize(fake_B0) self.fake_B1_resize = self.resize(fake_B1) + self.warped_fake_B0_2_resize = self.resize(warped_fake_B0_2) + self.warped_fake_B0_resize = self.resize(warped_fake_B0) + self.mutil_real_A0_tokens = self.netPreViT(self.real_A0_resize, self.atten_layers, get_tokens=True) self.mutil_real_A1_tokens = self.netPreViT(self.real_A1_resize, self.atten_layers, get_tokens=True) self.mutil_real_B0_tokens = self.netPreViT(real_B0, self.atten_layers, get_tokens=True) self.mutil_real_B1_tokens = self.netPreViT(real_B1, self.atten_layers, get_tokens=True) self.mutil_fake_B0_tokens = self.netPreViT(self.fake_B0_resize, self.atten_layers, get_tokens=True) self.mutil_fake_B1_tokens = self.netPreViT(self.fake_B1_resize, self.atten_layers, get_tokens=True) - - - if self.opt.phase == 'train': - # 真实图像的梯度 - real_gradient = torch.autograd.grad(self.real_B.sum(), self.real_B, create_graph=True)[0] - # 生成图像的梯度 - fake_gradient = torch.autograd.grad(self.fake_B.sum(), self.fake_B, create_graph=True)[0] - # 梯度图 - self.weight_real, self.weight_fake = self.cao.generate_weight_map(real_gradient, fake_gradient) - - # 生成图像的CTN光流图 - self.f_content = self.ctn(self.weight_fake) - - # 把前面生成后的图片再加上noisy_map - self.fake_B0_2 = self.fake_B0 + self.noisy_map - - # 变换后的图片 - wapped_fake_B0_2 = warp(self.fake_B0_2, self.f_content) - - # 经过第二次生成器 - self.fake_B0_2 = self.netG_2(wapped_fake_B0_2, self.time, z_in) - - - def tokens_concat(self, origin_tokens, adjacent_size): - adj_size = adjacent_size - B, token_num, C = origin_tokens.shape[0], origin_tokens.shape[1], origin_tokens.shape[2] - S = int(math.sqrt(token_num)) - if S * S != token_num: - print('Error! Not a square!') - token_map = origin_tokens.clone().reshape(B,S,S,C) - cut_patch_list = [] - for i in range(0, S, adj_size): - for j in range(0, S, adj_size): - i_left = i - i_right = i + adj_size + 1 if i + adj_size <= S else S + 1 - j_left = j - j_right = j + adj_size if j + adj_size <= S else S + 1 - - cut_patch = token_map[:, i_left:i_right, j_left: j_right, :] - cut_patch= cut_patch.reshape(B,-1,C) - cut_patch = torch.mean(cut_patch, dim=1, keepdim=True) - cut_patch_list.append(cut_patch) - - - result = torch.cat(cut_patch_list,dim=1) - return result - - - def cat_results(self, origin_tokens, adj_size_list): - res_list = [origin_tokens] - for ad_s in adj_size_list: - cat_result = self.tokens_concat(origin_tokens, ad_s) - res_list.append(cat_result) - - result = torch.cat(res_list, dim=1) - - return result + self.mutil_warped_fake_B0_tokens = self.netPreViT(self.warped_fake_B0_resize, self.atten_layers, get_tokens=True) + self.mutil_fake_B0_2_tokens = self.netPreViT(self.warped_fake_B0_2_resize, self.atten_layers, get_tokens=True) - def compute_D_loss(self): + def compute_D_loss(self): #判别器还是没有改 """Calculate GAN loss for the discriminator""" lambda_D_ViT = self.opt.lambda_D_ViT - fake_B0_tokens = self.mutil_fake_B0_tokens[self.opt.which_D_layer].detach() - fake_B1_tokens = self.mutil_fake_B1_tokens[self.opt.which_D_layer].detach() + fake_B0_tokens = self.mutil_fake_B0_tokens.detach() + fake_B1_tokens = self.mutil_fake_B1_tokens.detach() - real_B0_tokens = self.mutil_real_B0_tokens[self.opt.which_D_layer] - real_B1_tokens = self.mutil_real_B1_tokens[self.opt.which_D_layer] + real_B0_tokens = self.mutil_real_B0_tokens + real_B1_tokens = self.mutil_real_B1_tokens - fake_B0_tokens = self.cat_results(fake_B0_tokens, self.opt.adj_size_list) - fake_B1_tokens = self.cat_results(fake_B1_tokens, self.opt.adj_size_list) - - - - real_B0_tokens = self.cat_results(real_B0_tokens, self.opt.adj_size_list) - real_B1_tokens = self.cat_results(real_B1_tokens, self.opt.adj_size_list) - pre_fake0_ViT = self.netD_ViT(fake_B0_tokens) pre_fake1_ViT = self.netD_ViT(fake_B1_tokens) @@ -336,10 +290,9 @@ class ROMAModel(BaseModel): if self.opt.lambda_GAN > 0.0: - fake_B0_tokens = self.mutil_fake_B0_tokens[self.opt.which_D_layer] - fake_B1_tokens = self.mutil_fake_B1_tokens[self.opt.which_D_layer] - fake_B0_tokens = self.cat_results(fake_B0_tokens, self.opt.adj_size_list) - fake_B1_tokens = self.cat_results(fake_B1_tokens, self.opt.adj_size_list) + fake_B0_tokens = self.mutil_fake_B0_tokens + fake_B1_tokens = self.mutil_fake_B1_tokens + pred_fake0_ViT = self.netD_ViT(fake_B0_tokens) pred_fake1_ViT = self.netD_ViT(fake_B1_tokens) self.loss_G_GAN_ViT = (self.criterionGAN(pred_fake0_ViT, True) + self.criterionGAN(pred_fake1_ViT, True)) * 0.5 * self.opt.lambda_GAN @@ -357,8 +310,8 @@ class ROMAModel(BaseModel): # eq.9 ET_XY = self.netE(XtXt_1, self.time, XtXt_1).mean() - torch.logsumexp(self.netE(XtXt_1, self.time, XtXt_2).reshape(-1), dim=0) self.loss_SB = -(self.opt.num_timesteps - self.time[0]) / self.opt.num_timesteps * self.opt.tau * ET_XY - self.loss_SB += self.opt.tau * torch.mean((self.real_A_noisy - self.fake_B0) ** 2) - self.loss_SB += self.opt.tau * torch.mean((self.real_A_noisy2 - self.fake_B1) ** 2) + self.loss_SB += torch.mean((self.real_A_noisy - self.fake_B0) ** 2) + if self.opt.lambda_global > 0.0 or self.opt.lambda_spatial > 0.0: @@ -368,8 +321,8 @@ class ROMAModel(BaseModel): if self.opt.lambda_ctn > 0.0: - wapped_fake_B1 = warp(self.fake_B1, self.f_content) # use updated self.f_content - self.l2_loss = F.mse_loss(self.fake_B0_2, wapped_fake_B1) * self.opt.lambda_ctn + warped_fake_B1 = warp(self.fake_B0, self.f_content) # use updated self.f_content + self.l2_loss = F.mse_loss(self.warped_fake_B0_2, warped_fake_B1) * self.opt.lambda_ctn else: self.l2_loss = 0.0 diff --git a/models/self_build.py b/models/self_build.py index 31d5deb..952a097 100644 --- a/models/self_build.py +++ b/models/self_build.py @@ -79,37 +79,27 @@ class ContentAwareOptimization(nn.Module): cosine_sim = F.cosine_similarity(gradients, mean_grad, dim=2) # [B, N] return cosine_sim - def generate_weight_map(self, gradients_real, gradients_fake): + def generate_weight_map(self, gradients_fake): """ 生成内容感知权重图 Args: - gradients_real: [B, N, D] 真实图像判别器梯度 gradients_fake: [B, N, D] 生成图像判别器梯度 Returns: - weight_real: [B, N] 真实图像权重图 weight_fake: [B, N] 生成图像权重图 """ - # 计算真实图像块的余弦相似度 - cosine_real = self.compute_cosine_similarity(gradients_real) # [B, N] 公式5 # 计算生成图像块的余弦相似度 cosine_fake = self.compute_cosine_similarity(gradients_fake) # [B, N] # 选择内容丰富的区域（余弦相似度最低的eta_ratio比例） - k = int(self.eta_ratio * cosine_real.shape[1]) - - # 对真实图像生成权重图 - _, real_indices = torch.topk(-cosine_real, k, dim=1) # 选择最不相似的区域 - weight_real = torch.ones_like(cosine_real) - for b in range(cosine_real.shape[0]): - weight_real[b, real_indices[b]] = self.lambda_inc / (1e-6 + torch.abs(cosine_real[b, real_indices[b]])) #公式6 - + k = int(self.eta_ratio * cosine_fake.shape[1]) + # 对生成图像生成权重图（同理） _, fake_indices = torch.topk(-cosine_fake, k, dim=1) weight_fake = torch.ones_like(cosine_fake) for b in range(cosine_fake.shape[0]): weight_fake[b, fake_indices[b]] = self.lambda_inc / (1e-6 + torch.abs(cosine_fake[b, fake_indices[b]])) - return weight_real, weight_fake + return weight_fake def forward(self, D_real, D_fake, real_scores, fake_scores): """ @@ -458,7 +448,7 @@ class CTNxModel(BaseModel): self.real_A_noisy = Xt.detach() self.real_A_noisy2 = Xt2.detach() # 保存noisy_map - self.noisy_map = self.real_A_noisy - self.real_A + self.noisy_map = self.real_A_noisy - self.real_A0 # ============ 第三步：拼接输入并执行网络推理 ============= bs = self.mutil_real_A0_tokens.size(0)