import numpy as np import torch from .base_model import BaseModel from . import networks from .patchnce import PatchNCELoss import util.util as util import timm import time import torch.nn.functional as F import sys from functools import partial import torch.nn as nn import math from torchvision.transforms import transforms as tfs class ROMAModel(BaseModel): @staticmethod def modify_commandline_options(parser, is_train=True): """ Configures options specific for CUT model """ parser.add_argument('--adj_size_list', type=list, default=[2, 4, 6, 8, 12], help='different scales of perception field') parser.add_argument('--lambda_mlp', type=float, default=1.0, help='weight of lr for discriminator') parser.add_argument('--lambda_motion', type=float, default=1.0, help='weight for Temporal Consistency') parser.add_argument('--lambda_D_ViT', type=float, default=1.0, help='weight for discriminator') parser.add_argument('--lambda_GAN', type=float, default=1.0, help='weight for GAN loss: GAN(G(X))') parser.add_argument('--lambda_global', type=float, default=1.0, help='weight for Global Structural Consistency') parser.add_argument('--lambda_spatial', type=float, default=1.0, help='weight for Local Structural Consistency') parser.add_argument('--atten_layers', type=str, default='1,3,5', help='compute Cross-Similarity on which layers') parser.add_argument('--local_nums', type=int, default=256) parser.add_argument('--which_D_layer', type=int, default=-1) parser.add_argument('--side_length', type=int, default=7) parser.set_defaults(pool_size=0) opt, _ = parser.parse_known_args() return parser def __init__(self, opt): BaseModel.__init__(self, opt) self.loss_names = ['G_GAN_ViT', 'D_real_ViT', 'D_fake_ViT', 'global', 'spatial', 'motion'] self.visual_names = ['real_A0', 'real_A1', 'fake_B0', 'fake_B1', 'real_B0', 'real_B1'] self.atten_layers = [int(i) for i in self.opt.atten_layers.split(',')] if self.isTrain: self.model_names = ['G', 'D_ViT'] else: # during test time, only load G self.model_names = ['G'] # define networks (both generator and discriminator) self.netG = 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) if self.isTrain: self.netD_ViT = networks.MLPDiscriminator().to(self.device) self.netPreViT = timm.create_model("vit_base_patch16_384",pretrained=True).to(self.device) self.norm = F.softmax self.resize = tfs.Resize(size=(384,384)) self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device) self.criterionNCE = [] for atten_layer in self.atten_layers: self.criterionNCE.append(PatchNCELoss(opt).to(self.device)) self.criterionL1 = torch.nn.L1Loss().to(self.device) self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2)) self.optimizer_D_ViT = torch.optim.Adam(self.netD_ViT.parameters(), lr=opt.lr * opt.lambda_mlp, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D_ViT) def data_dependent_initialize(self, data): """ The feature network netF is defined in terms of the shape of the intermediate, extracted features of the encoder portion of netG. Because of this, the weights of netF are initialized at the first feedforward pass with some input images. Please also see PatchSampleF.create_mlp(), which is called at the first forward() call. """ pass def optimize_parameters(self): # forward self.forward() # update D self.set_requires_grad(self.netD_ViT, True) self.optimizer_D_ViT.zero_grad() self.loss_D = self.compute_D_loss() self.loss_D.backward() self.optimizer_D_ViT.step() # update G self.set_requires_grad(self.netD_ViT, False) self.optimizer_G.zero_grad() self.loss_G = self.compute_G_loss() self.loss_G.backward() self.optimizer_G.step() def set_input(self, input): """Unpack input data from the dataloader and perform necessary pre-processing steps. Parameters: input (dict): include the data itself and its metadata information. The option 'direction' can be used to swap domain A and domain B. """ AtoB = self.opt.direction == 'AtoB' self.real_A0 = input['A0' if AtoB else 'B0'].to(self.device) self.real_A1 = input['A1' if AtoB else 'B1'].to(self.device) self.real_B0 = input['B0' if AtoB else 'A0'].to(self.device) self.real_B1 = input['B1' if AtoB else 'A1'].to(self.device) self.image_paths = input['A_paths' if AtoB else 'B_paths'] def forward(self): """Run forward pass; called by both functions and .""" # ============ 第一步：对 real_A / real_A2 进行多步随机生成过程 ============ tau = self.opt.tau T = self.opt.num_timesteps incs = np.array([0] + [1/(i+1) for i in range(T-1)]) times = np.cumsum(incs) times = times / times[-1] times = 0.5 * times[-1] + 0.5 * times #[0.5,1] times = np.concatenate([np.zeros(1), times]) times = torch.tensor(times).float().cuda() self.times = times bs = self.mutil_real_A0_tokens.size(0) time_idx = (torch.randint(T, size=[1]).cuda() * torch.ones(size=[1]).cuda()).long() self.time_idx = time_idx with torch.no_grad(): self.netG.eval() # ============ 第二步：对 real_A / real_A2 进行多步随机生成过程 ============ for t in range(self.time_idx.int().item() + 1): # 计算增量 delta 与 inter/scale，用于每个时间步的插值等 if t > 0: delta = times[t] - times[t - 1] denom = times[-1] - times[t - 1] inter = (delta / denom).reshape(-1, 1, 1, 1) scale = (delta * (1 - delta / denom)).reshape(-1, 1, 1, 1) # 对 Xt、Xt2 进行随机噪声更新 Xt = self.mutil_real_A0_tokens if (t == 0) else (1 - inter) * Xt + inter * Xt_1.detach() + \ (scale * tau).sqrt() * torch.randn_like(Xt).to(self.mutil_real_A0_tokens.device) time_idx = (t * torch.ones(size=[self.mutil_real_A0_tokens.shape[0]]).to(self.mutil_real_A0_tokens.device)).long() z = torch.randn(size=[self.mutil_real_A0_tokens.shape[0], 4 * self.opt.ngf]).to(self.mutil_real_A0_tokens.device) self.time = times[time_idx] Xt_1 = self.netG(Xt, self.time, z) Xt2 = self.mutil_real_A1_tokens if (t == 0) else (1 - inter) * Xt2 + inter * Xt_12.detach() + \ (scale * tau).sqrt() * torch.randn_like(Xt2).to(self.mutil_real_A1_tokens.device) time_idx = (t * torch.ones(size=[self.mutil_real_A1_tokens.shape[0]]).to(self.mutil_real_A1_tokens.device)).long() z = torch.randn(size=[self.mutil_real_A1_tokens.shape[0], 4 * self.opt.ngf]).to(self.mutil_real_A1_tokens.device) Xt_12 = self.netG(Xt2, self.time, z) # 保存去噪后的中间结果 (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.mutil_real_A0_tokens.size(0) z_in = torch.randn(size=[2 * bs, 4 * self.opt.ngf]).to(self.mutil_real_A0_tokens.device) z_in2 = torch.randn(size=[bs, 4 * self.opt.ngf]).to(self.mutil_real_A1_tokens.device) # 将 real_A, real_B 拼接 (如 nce_idt=True)，并同样处理 real_A_noisy 与 XtB self.real = self.mutil_real_A0_tokens self.realt = self.real_A_noisy if self.opt.flip_equivariance: self.flipped_for_equivariance = self.opt.isTrain and (np.random.random() < 0.5) if self.flipped_for_equivariance: self.real = torch.flip(self.real, [3]) self.realt = torch.flip(self.realt, [3]) self.fake_B0 = self.netG(self.real_A0) self.fake_B1 = self.netG(self.real_A1) if self.opt.isTrain: real_A0 = self.real_A0 real_A1 = self.real_A1 real_B0 = self.real_B0 real_B1 = self.real_B1 fake_B0 = self.fake_B0 fake_B1 = self.fake_B1 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.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) 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 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() 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] 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) self.loss_D_fake_ViT = (self.criterionGAN(pre_fake0_ViT, False).mean() + self.criterionGAN(pre_fake1_ViT, False).mean()) * 0.5 * lambda_D_ViT pred_real0_ViT = self.netD_ViT(real_B0_tokens) pred_real1_ViT = self.netD_ViT(real_B1_tokens) self.loss_D_real_ViT = (self.criterionGAN(pred_real0_ViT, True).mean() + self.criterionGAN(pred_real1_ViT, True).mean()) * 0.5 * lambda_D_ViT self.loss_D_ViT = (self.loss_D_fake_ViT + self.loss_D_real_ViT) * 0.5 return self.loss_D_ViT def compute_G_loss(self): 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) 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 else: self.loss_G_GAN_ViT = 0.0 if self.opt.lambda_global > 0.0 or self.opt.lambda_spatial > 0.0: self.loss_global, self.loss_spatial = self.calculate_attention_loss() else: self.loss_global, self.loss_spatial = 0.0, 0.0 if self.opt.lambda_motion > 0.0: self.loss_motion = 0.0 for real_A0_tokens, real_A1_tokens, fake_B0_tokens, fake_B1_tokens in zip(self.mutil_real_A0_tokens, self.mutil_real_A1_tokens, self.mutil_fake_B0_tokens, self.mutil_fake_B1_tokens): A0_B1 = real_A0_tokens.bmm(fake_B1_tokens.permute(0,2,1)) B0_A1 = fake_B0_tokens.bmm(real_A1_tokens.permute(0,2,1)) cos_dis_global = F.cosine_similarity(A0_B1, B0_A1, dim=-1) self.loss_motion += self.criterionL1(torch.ones_like(cos_dis_global), cos_dis_global).mean() else: self.loss_motion = 0.0 self.loss_G = self.loss_G_GAN_ViT + self.loss_global + self.loss_spatial + self.loss_motion return self.loss_G def calculate_attention_loss(self): n_layers = len(self.atten_layers) mutil_real_A0_tokens = self.mutil_real_A0_tokens mutil_real_A1_tokens = self.mutil_real_A1_tokens mutil_fake_B0_tokens = self.mutil_fake_B0_tokens mutil_fake_B1_tokens = self.mutil_fake_B1_tokens if self.opt.lambda_global > 0.0: loss_global = self.calculate_similarity(mutil_real_A0_tokens, mutil_fake_B0_tokens) + self.calculate_similarity(mutil_real_A1_tokens, mutil_fake_B1_tokens) loss_global *= 0.5 else: loss_global = 0.0 if self.opt.lambda_spatial > 0.0: loss_spatial = 0.0 local_nums = self.opt.local_nums tokens_cnt = 576 local_id = np.random.permutation(tokens_cnt) local_id = local_id[:int(min(local_nums, tokens_cnt))] mutil_real_A0_local_tokens = self.netPreViT(self.real_A0_resize, self.atten_layers, get_tokens=True, local_id=local_id, side_length = self.opt.side_length) mutil_real_A1_local_tokens = self.netPreViT(self.real_A1_resize, self.atten_layers, get_tokens=True, local_id=local_id, side_length = self.opt.side_length) mutil_fake_B0_local_tokens = self.netPreViT(self.fake_B0_resize, self.atten_layers, get_tokens=True, local_id=local_id, side_length = self.opt.side_length) mutil_fake_B1_local_tokens = self.netPreViT(self.fake_B1_resize, self.atten_layers, get_tokens=True, local_id=local_id, side_length = self.opt.side_length) loss_spatial = self.calculate_similarity(mutil_real_A0_local_tokens, mutil_fake_B0_local_tokens) + self.calculate_similarity(mutil_real_A1_local_tokens, mutil_fake_B1_local_tokens) loss_spatial *= 0.5 else: loss_spatial = 0.0 return loss_global * self.opt.lambda_global, loss_spatial * self.opt.lambda_spatial def calculate_similarity(self, mutil_src_tokens, mutil_tgt_tokens): loss = 0.0 n_layers = len(self.atten_layers) for src_tokens, tgt_tokens in zip(mutil_src_tokens, mutil_tgt_tokens): src_tgt = src_tokens.bmm(tgt_tokens.permute(0,2,1)) tgt_src = tgt_tokens.bmm(src_tokens.permute(0,2,1)) cos_dis_global = F.cosine_similarity(src_tgt, tgt_src, dim=-1) loss += self.criterionL1(torch.ones_like(cos_dis_global), cos_dis_global).mean() loss = loss / n_layers return loss