385 lines
17 KiB
Python
385 lines
17 KiB
Python
import numpy as np
|
||
import torch
|
||
from .base_model import BaseModel
|
||
from . import networks
|
||
from .patchnce import PatchNCELoss
|
||
from .cnt import *
|
||
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 ROMAUNSBModel(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_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('--lambda_inc', type=float, default=2.0, help='weight for Content Aware Optimization')
|
||
parser.add_argument('--eta_ratio', type=float, default=0.4, help='ratio for selecting content region')
|
||
|
||
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']
|
||
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', 'G_2']
|
||
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)
|
||
|
||
# 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)
|
||
|
||
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.optimizer_E = torch.optim.Adam(self.netE.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)
|
||
self.optimizers.append(self.optimizer_E)
|
||
|
||
self.cao = ContentAwareOptimization(opt.lambda_inc, opt.eta_ratio) #损失函数
|
||
self.ctn = ContentAwareTemporalNorm() #生成的伪光流场
|
||
|
||
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 E
|
||
self.set_requires_grad(self.netE, True)
|
||
self.optimizer_E.zero_grad()
|
||
self.loss_E = self.compute_E_loss()
|
||
self.loss_E.backward()
|
||
self.optimizer_E.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 <optimize_parameters> and <test>."""
|
||
|
||
# ============ 第一步：对 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.real_A0.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.real_A0 if (t == 0) else (1 - inter) * Xt + inter * Xt_1.detach() + \
|
||
(scale * tau).sqrt() * torch.randn_like(Xt).to(self.real_A0.device)
|
||
time_idx = (t * torch.ones(size=[self.real_A0.shape[0]]).to(self.real_A0.device)).long()
|
||
z = torch.randn(size=[self.real_A0.shape[0], 4 * self.opt.ngf]).to(self.real_A0.device)
|
||
self.time = times[time_idx]
|
||
Xt_1 = self.netG(Xt, self.time, z)
|
||
|
||
Xt2 = self.real_A1 if (t == 0) else (1 - inter) * Xt2 + inter * Xt_12.detach() + \
|
||
(scale * tau).sqrt() * torch.randn_like(Xt2).to(self.real_A1.device)
|
||
time_idx = (t * torch.ones(size=[self.real_A1.shape[0]]).to(self.real_A1.device)).long()
|
||
z = torch.randn(size=[self.real_A1.shape[0], 4 * self.opt.ngf]).to(self.real_A1.device)
|
||
Xt_12 = self.netG(Xt2, self.time, z)
|
||
|
||
# 保存去噪后的中间结果 (real_A_noisy 等)，供下一步做拼接
|
||
self.real_A_noisy = Xt.detach()
|
||
self.real_A_noisy2 = Xt2.detach()
|
||
|
||
# ============ 第三步：拼接输入并执行网络推理 =============
|
||
bs = self.real_A0.size(0)
|
||
z_in = torch.randn(size=[2 * bs, 4 * self.opt.ngf]).to(self.real_A0.device)
|
||
z_in2 = torch.randn(size=[bs, 4 * self.opt.ngf]).to(self.real_A1.device)
|
||
# 将 real_A, real_B 拼接 (如 nce_idt=True)，并同样处理 real_A_noisy 与 XtB
|
||
self.real = self.real_A0
|
||
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.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
|
||
real_B0 = self.real_B0
|
||
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)
|
||
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): #判别器还是没有改
|
||
"""Calculate GAN loss for the discriminator"""
|
||
|
||
|
||
lambda_D_ViT = self.opt.lambda_D_ViT
|
||
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
|
||
real_B1_tokens = self.mutil_real_B1_tokens
|
||
|
||
|
||
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_E_loss(self):
|
||
"""计算判别器 E 的损失"""
|
||
|
||
XtXt_1 = torch.cat([self.real_A_noisy, self.fake_B0.detach()], dim=1)
|
||
XtXt_2 = torch.cat([self.real_A_noisy2, self.fake_B1.detach()], dim=1)
|
||
temp = torch.logsumexp(self.netE(XtXt_1, self.time, XtXt_2).reshape(-1), dim=0).mean()
|
||
self.loss_E = -self.netE(XtXt_1, self.time, XtXt_1).mean() + temp + temp**2
|
||
|
||
return self.loss_E
|
||
|
||
|
||
def compute_G_loss(self):
|
||
|
||
if self.opt.lambda_GAN > 0.0:
|
||
|
||
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
|
||
else:
|
||
self.loss_G_GAN_ViT = 0.0
|
||
|
||
|
||
self.loss_SB = 0
|
||
if self.opt.lambda_SB > 0.0:
|
||
XtXt_1 = torch.cat([self.real_A_noisy, self.fake_B0], dim=1)
|
||
XtXt_2 = torch.cat([self.real_A_noisy2, self.fake_B1], dim=1)
|
||
|
||
bs = self.opt.batch_size
|
||
|
||
# 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 += torch.mean((self.real_A_noisy - self.fake_B0) ** 2)
|
||
|
||
|
||
|
||
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_ctn > 0.0:
|
||
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
|
||
|
||
|
||
self.loss_G = self.loss_G_GAN_ViT + self.loss_global + self.loss_spatial + self.l2_loss # include l2_loss in total loss
|
||
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
|
||
|