├── models ├── __init__.py ├── DiffAug.py ├── VITGen.py ├── QEM.py ├── VIT_old.py ├── PSM.py ├── CNNDis.py ├── VIT.py └── ops.py ├── util ├── __init__.py ├── data_prefetcher.py ├── inception.py └── misc.py ├── assets └── demo.jpg ├── losses ├── reconstruct.py ├── __init__.py └── perceptual.py ├── datasets.py ├── evaluate.py ├── README.md ├── main.py ├── engine.py └── LICENSE /models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /util/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /assets/demo.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kaiseem/QueryOTR/HEAD/assets/demo.jpg -------------------------------------------------------------------------------- /losses/reconstruct.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from einops import rearrange 4 | 5 | class ReconLoss(nn.Module): 6 | def __init__(self,image_size=192,crop_width=32,loss_type='mse'): 7 | super(ReconLoss, self).__init__() 8 | assert loss_type in ['l1','mse'] 9 | mask = torch.ones((image_size, image_size)) 10 | mask[crop_width:-crop_width, crop_width:-crop_width] = 0 11 | self.mask=mask.view(-1).long().cuda() 12 | self.outer_index=self.mask==1 13 | if loss_type=='l1': 14 | self.loss=nn.L1Loss() 15 | else: 16 | self.loss=nn.MSELoss() 17 | 18 | def forward(self,input_fake, input_real): 19 | input_fake = rearrange(input_fake, 'b c w h -> b (w h) c')[:, self.outer_index] 20 | input_real = rearrange(input_real, 'b c w h -> b (w h) c')[:, self.outer_index] 21 | return self.loss(input_fake,input_real) 22 | 23 | -------------------------------------------------------------------------------- /util/data_prefetcher.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | class data_prefetcher(): 4 | def __init__(self, loader, device, prefetch=True): 5 | self.loader = iter(loader) 6 | self.prefetch = prefetch 7 | self.device = device 8 | if prefetch: 9 | self.stream = torch.cuda.Stream() 10 | self.preload() 11 | 12 | def preload(self): 13 | try: 14 | self.next_samples = next(self.loader) 15 | except StopIteration: 16 | self.next_samples = None 17 | return 18 | with torch.cuda.stream(self.stream): 19 | for k, v in self.next_samples.items(): 20 | self.next_samples[k]=v.to(self.device) 21 | 22 | def next(self): 23 | if self.prefetch: 24 | torch.cuda.current_stream().wait_stream(self.stream) 25 | samples = self.next_samples 26 | if samples is not None: 27 | for k, v in samples.items(): 28 | v.record_stream(torch.cuda.current_stream()) 29 | self.preload() 30 | else: 31 | try: 32 | samples = next(self.loader) 33 | for k, v in self.next_samples.items(): 34 | self.next_samples[k] = v.to(self.device) 35 | except StopIteration: 36 | samples = None 37 | return samples 38 | 39 | 40 | -------------------------------------------------------------------------------- /losses/__init__.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from .reconstruct import ReconLoss 3 | from .perceptual import PerceptualLoss 4 | import torch 5 | from torchvision import transforms 6 | from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD 7 | 8 | class SetCriterion(nn.Module): 9 | def __init__(self,opts): 10 | super().__init__() 11 | self.recon_loss=ReconLoss() 12 | self.perceptual_loss=PerceptualLoss() 13 | self.gen_weight_dict={'loss_g_recon':5, 'loss_g_adversarial':1, 'loss_g_perceptual':10} 14 | self.dis_weight_dict = {'loss_d_adversarial': 1} 15 | self.imagenet_normalize=transforms.Normalize( mean=torch.tensor(IMAGENET_DEFAULT_MEAN), std=torch.tensor(IMAGENET_DEFAULT_STD)) 16 | self.patch_mean=opts.patch_mean 17 | self.patch_std=opts.patch_std 18 | 19 | def renorm(self,tensor): 20 | tensor = tensor * self.patch_std + self.patch_mean 21 | return self.imagenet_normalize(tensor) 22 | 23 | def get_dis_loss(self, input_fake, input_real, discriminator=None): 24 | assert discriminator is not None 25 | return {'loss_d_adversarial': discriminator.calc_dis_loss(input_fake.detach(),input_real)} 26 | 27 | def get_gen_loss(self, input_fake, input_real,discriminator=None, warmup=False): 28 | if not warmup: 29 | assert discriminator is not None 30 | g_loss_dict={'loss_g_adversarial': discriminator.calc_gen_loss(input_fake,input_real)} 31 | g_loss_dict['loss_g_recon']=self.recon_loss(input_fake,input_real) 32 | g_loss_dict['loss_g_perceptual']=self.perceptual_loss(self.renorm(input_fake),self.renorm(input_real)) 33 | return g_loss_dict 34 | else: 35 | return {'loss_g_recon':self.recon_loss(input_fake,input_real)} 36 | 37 | 38 | 39 | 40 | 41 | -------------------------------------------------------------------------------- /losses/perceptual.py: -------------------------------------------------------------------------------- 1 | from torchvision.models.vgg import vgg19 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | 7 | class VGG19(torch.nn.Module): 8 | def __init__(self, requires_grad=False): 9 | super().__init__() 10 | vgg_pretrained_features = vgg19(pretrained=True).features 11 | self.slice1 = torch.nn.Sequential() 12 | self.slice2 = torch.nn.Sequential() 13 | self.slice3 = torch.nn.Sequential() 14 | self.slice4 = torch.nn.Sequential() 15 | self.slice5 = torch.nn.Sequential() 16 | for x in range(2): 17 | self.slice1.add_module(str(x), vgg_pretrained_features[x]) 18 | for x in range(2, 7): 19 | self.slice2.add_module(str(x), vgg_pretrained_features[x]) 20 | for x in range(7, 12): 21 | self.slice3.add_module(str(x), vgg_pretrained_features[x]) 22 | for x in range(12, 21): 23 | self.slice4.add_module(str(x), vgg_pretrained_features[x]) 24 | for x in range(21, 30): 25 | self.slice5.add_module(str(x), vgg_pretrained_features[x]) 26 | if not requires_grad: 27 | for param in self.parameters(): 28 | param.requires_grad = False 29 | 30 | def forward(self, X): 31 | h_relu1 = self.slice1(X) 32 | h_relu2 = self.slice2(h_relu1) 33 | h_relu3 = self.slice3(h_relu2) 34 | h_relu4 = self.slice4(h_relu3) 35 | h_relu5 = self.slice5(h_relu4) 36 | out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5] 37 | return out 38 | 39 | 40 | class PerceptualLoss(nn.Module): 41 | ''' 42 | same as https://github.com/NVlabs/SPADE/blob/master/models/networks/loss.py 43 | ''' 44 | def __init__(self): 45 | super(PerceptualLoss, self).__init__() 46 | 47 | self.vgg = VGG19().cuda() 48 | self.criterion = nn.L1Loss() 49 | self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0] 50 | 51 | def forward(self, x, y): 52 | x_vgg, y_vgg = self.vgg(x), self.vgg(y) 53 | loss = 0 54 | for i in range(len(x_vgg)): 55 | loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach()) 56 | return loss 57 | 58 | # self.vgg_model= vgg16(pretrained=True) 59 | # self.vgg_model.features.__delitem__(23)#delete maxpool 60 | # self.vgg_model.features.__delitem__(-1)#delete maxpool 61 | # self.instance_norm=nn.InstanceNorm2d(512,affine=False) 62 | # self.vgg_model.eval() 63 | # for param in self.vgg_model.parameters(): 64 | # param.requires_grad = False 65 | # self.cuda() 66 | # 67 | # def forward(self, input_fake, input_real): 68 | # return torch.mean((self.instance_norm(self.vgg_model.features(input_fake)) - self.instance_norm(self.vgg_model.features(input_real))) ** 2) 69 | 70 | 71 | -------------------------------------------------------------------------------- /models/DiffAug.py: -------------------------------------------------------------------------------- 1 | # Differentiable Augmentation for Data-Efficient GAN Training 2 | # Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han 3 | # https://arxiv.org/pdf/2006.10738 4 | 5 | import torch 6 | import torch.nn.functional as F 7 | 8 | def DiffAugment(x, policy='', channels_first=True): 9 | if policy: 10 | if not channels_first: 11 | x = x.permute(0, 3, 1, 2) 12 | for p in policy.split(','): 13 | for f in AUGMENT_FNS[p]: 14 | x = f(x) 15 | if not channels_first: 16 | x = x.permute(0, 2, 3, 1) 17 | x = x.contiguous() 18 | return x 19 | 20 | 21 | def rand_brightness(x): 22 | x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5) 23 | return x 24 | 25 | 26 | def rand_saturation(x): 27 | x_mean = x.mean(dim=1, keepdim=True) 28 | x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean 29 | return x 30 | 31 | 32 | def rand_contrast(x): 33 | x_mean = x.mean(dim=[1, 2, 3], keepdim=True) 34 | x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean 35 | return x 36 | 37 | 38 | def rand_translation(x, ratio=0.125): 39 | shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5) 40 | translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device) 41 | translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device) 42 | grid_batch, grid_x, grid_y = torch.meshgrid( 43 | torch.arange(x.size(0), dtype=torch.long, device=x.device), 44 | torch.arange(x.size(2), dtype=torch.long, device=x.device), 45 | torch.arange(x.size(3), dtype=torch.long, device=x.device), 46 | ) 47 | grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1) 48 | grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1) 49 | x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0]) 50 | x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2).contiguous() 51 | return x 52 | 53 | 54 | def rand_cutout(x, ratio=0.5): 55 | cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5) 56 | offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device) 57 | offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device) 58 | grid_batch, grid_x, grid_y = torch.meshgrid( 59 | torch.arange(x.size(0), dtype=torch.long, device=x.device), 60 | torch.arange(cutout_size[0], dtype=torch.long, device=x.device), 61 | torch.arange(cutout_size[1], dtype=torch.long, device=x.device), 62 | ) 63 | grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1) 64 | grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1) 65 | mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device) 66 | mask[grid_batch, grid_x, grid_y] = 0 67 | x = x * mask.unsqueeze(1) 68 | return x 69 | 70 | 71 | AUGMENT_FNS = { 72 | 'color': [rand_brightness, rand_saturation, rand_contrast], 73 | 'translation': [rand_translation], 74 | 'cutout': [rand_cutout], 75 | } -------------------------------------------------------------------------------- /datasets.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 4 | 5 | from PIL import Image 6 | import torch.utils.data as data 7 | import torch 8 | from torchvision import transforms 9 | from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD 10 | 11 | 12 | IMG_EXTENSIONS = [ 13 | '.jpg', '.JPG', '.jpeg', '.JPEG', 14 | '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', 15 | '.tif', '.TIF', '.tiff', '.TIFF','npy','mat' 16 | ] 17 | def is_image_file(filename): 18 | return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) 19 | 20 | def make_dataset(dir, max_dataset_size=float("inf")): 21 | images = [] 22 | assert os.path.isdir(dir), '%s is not a valid directory' % dir 23 | for root, _, fnames in sorted(os.walk(dir)): 24 | for fname in fnames: 25 | if is_image_file(fname): 26 | path = os.path.join(root, fname) 27 | images.append(path) 28 | return images[:min(max_dataset_size, len(images))] 29 | from copy import deepcopy 30 | 31 | class ImageDataset(data.Dataset): 32 | def __init__(self, opts): 33 | self.img_paths = sorted(make_dataset(opts.data_root)) 34 | self.is_train=not opts.eval 35 | input_size=opts.input_size 36 | output_size=opts.output_size 37 | per_edge_pad=(output_size-input_size)//2 38 | normlize_target=opts.normlize_target 39 | patch_mean=opts.patch_mean 40 | patch_std=opts.patch_std 41 | 42 | if self.is_train: 43 | self.transform = transforms.Compose([ 44 | transforms.RandomResizedCrop(output_size), 45 | transforms.RandomHorizontalFlip(), 46 | transforms.ToTensor(), 47 | ]) 48 | else: 49 | self.transform = transforms.Compose([ 50 | transforms.Resize((output_size, output_size)), 51 | transforms.ToTensor(), 52 | ]) 53 | 54 | self.input_image_normalize=transforms.Normalize( mean=torch.tensor(IMAGENET_DEFAULT_MEAN), std=torch.tensor(IMAGENET_DEFAULT_STD)) 55 | if normlize_target: 56 | self.output_patch_normalize=transforms.Normalize( mean=torch.tensor((patch_mean,patch_mean,patch_mean)), std=torch.tensor((patch_std,patch_std,patch_std))) 57 | else: 58 | self.output_patch_normalize=self.input_image_normalize 59 | 60 | self._mean=torch.tensor((patch_mean,patch_mean,patch_mean)) 61 | self._std=torch.tensor((patch_std,patch_std,patch_std)) 62 | 63 | self.mask=torch.zeros([1,output_size, output_size]) 64 | self.mask[:,per_edge_pad:-per_edge_pad,per_edge_pad:-per_edge_pad]=1 65 | 66 | self.per_edge_pad=per_edge_pad 67 | 68 | def __getitem__(self, index): 69 | name= os.path.splitext(os.path.split(self.img_paths[index])[-1])[0] 70 | im=Image.open(self.img_paths[index]).convert('RGB') 71 | im=self.transform(im) 72 | input_img=self.input_image_normalize(deepcopy(im))[:,self.per_edge_pad:-self.per_edge_pad,self.per_edge_pad:-self.per_edge_pad] 73 | 74 | gt=self.output_patch_normalize(deepcopy(im)) 75 | gt_inner=deepcopy(gt)*self.mask 76 | return {'input':input_img,'ground_truth':gt,'gt_inner':gt_inner,'name':name} 77 | 78 | def __len__(self): 79 | return len(self.img_paths) 80 | -------------------------------------------------------------------------------- /util/inception.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.autograd import Variable 4 | from torch.nn import functional as F 5 | import torch.utils.data 6 | 7 | from torchvision.models.inception import inception_v3 8 | 9 | import numpy as np 10 | from scipy.stats import entropy 11 | 12 | def inception_score(imgs, cuda=True, batch_size=32, resize=False, splits=1): 13 | """Computes the inception score of the generated images imgs 14 | imgs -- Torch dataset of (3xHxW) numpy images normalized in the range [-1, 1] 15 | cuda -- whether or not to run on GPU 16 | batch_size -- batch size for feeding into Inception v3 17 | splits -- number of splits 18 | """ 19 | N = len(imgs) 20 | 21 | assert batch_size > 0 22 | assert N > batch_size 23 | 24 | # Set up dtype 25 | if cuda: 26 | dtype = torch.cuda.FloatTensor 27 | else: 28 | if torch.cuda.is_available(): 29 | print("WARNING: You have a CUDA device, so you should probably set cuda=True") 30 | dtype = torch.FloatTensor 31 | 32 | # Set up dataloader 33 | dataloader = torch.utils.data.DataLoader(imgs, batch_size=batch_size) 34 | 35 | # Load inception model 36 | inception_model = inception_v3(pretrained=True, transform_input=False).type(dtype) 37 | inception_model.eval(); 38 | up = nn.Upsample(size=(299, 299), mode='bilinear').type(dtype) 39 | def get_pred(x): 40 | if resize: 41 | x = up(x) 42 | x = inception_model(x) 43 | return F.softmax(x).data.cpu().numpy() 44 | 45 | # Get predictions 46 | preds = np.zeros((N, 1000)) 47 | 48 | for i, batch in enumerate(dataloader, 0): 49 | batch = batch.type(dtype) 50 | batchv = Variable(batch) 51 | batch_size_i = batch.size()[0] 52 | 53 | preds[i*batch_size:i*batch_size + batch_size_i] = get_pred(batchv) 54 | 55 | # Now compute the mean kl-div 56 | split_scores = [] 57 | 58 | for k in range(splits): 59 | part = preds[k * (N // splits): (k+1) * (N // splits), :] 60 | py = np.mean(part, axis=0) 61 | scores = [] 62 | for i in range(part.shape[0]): 63 | pyx = part[i, :] 64 | scores.append(entropy(pyx, py)) 65 | split_scores.append(np.exp(np.mean(scores))) 66 | 67 | return np.mean(split_scores), np.std(split_scores) 68 | 69 | if __name__ == '__main__': 70 | class IgnoreLabelDataset(torch.utils.data.Dataset): 71 | def __init__(self, orig): 72 | self.orig = orig 73 | 74 | def __getitem__(self, index): 75 | return self.orig[index][0] 76 | 77 | def __len__(self): 78 | return len(self.orig) 79 | 80 | import torchvision.datasets as dset 81 | import torchvision.transforms as transforms 82 | 83 | cifar = dset.CIFAR10(root='data/', download=True, 84 | transform=transforms.Compose([ 85 | transforms.Scale(32), 86 | transforms.ToTensor(), 87 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) 88 | ]) 89 | ) 90 | 91 | IgnoreLabelDataset(cifar) 92 | 93 | print ("Calculating Inception Score...") 94 | print (inception_score(IgnoreLabelDataset(cifar), cuda=True, batch_size=32, resize=True, splits=10)) -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import os 2 | from einops import rearrange 3 | import matplotlib.pyplot as plt 4 | os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' 5 | from torch.utils.data import Dataset 6 | from models.VITGen import TransGen 7 | from datasets import ImageDataset 8 | import argparse 9 | from util.inception import inception_score 10 | import numpy as np 11 | from PIL import Image 12 | import torch 13 | parser = argparse.ArgumentParser() 14 | parser.add_argument('--eval', default=True, type=bool) 15 | parser.add_argument('--batch_size', type=int, default=64) 16 | parser.add_argument("-r", "--resume", type=str) 17 | parser.add_argument('--input_size', type=int, default=128) 18 | parser.add_argument('--output_size', type=int, default=192) 19 | parser.add_argument('--dec_depth', type=int, default=4) 20 | parser.add_argument('--normlize_target', default=True, type=bool, help='normalized the target patch pixels') 21 | parser.add_argument('--patch_mean', type=float, default=0.5044838) 22 | parser.add_argument('--patch_std', type=float, default=0.1355051) 23 | parser.add_argument('--data_root', type=str, default='') 24 | parser.add_argument('--epoch', type=str, default=None) 25 | opts = parser.parse_args() 26 | 27 | 28 | 29 | def denorm_img(tensor): 30 | _mean = torch.tensor([opts.patch_mean, opts.patch_mean, opts.patch_mean]).unsqueeze(-1).unsqueeze(-1).unsqueeze(0) 31 | _std = torch.tensor([opts.patch_std, opts.patch_std, opts.patch_std]).unsqueeze(-1).unsqueeze(-1).unsqueeze(0) 32 | tensor = tensor * _std.cuda().expand_as(tensor) + _mean.cuda().expand_as(tensor) 33 | tensor = rearrange(tensor[0:1], 'b c w h -> b w h c').detach().cpu() 34 | tensor = np.clip(tensor[0].numpy(), 0, 1) 35 | return tensor 36 | 37 | if __name__=='__main__': 38 | gen = TransGen(opts=opts).cuda() 39 | logdir = opts.resume 40 | ckptdir = os.path.join(logdir, "checkpoints") 41 | if opts.epoch is not None: 42 | ckpt_name = os.path.join(ckptdir, f'{opts.epoch}.pth') 43 | else: 44 | ckpt_name = os.path.join(ckptdir, 'latest.pth') 45 | assert os.path.isfile(ckpt_name), f'check if existing checkpoint files {ckpt_name}' 46 | gtdir = os.path.join(logdir, "gt") 47 | generatedir = os.path.join(logdir, "generated") 48 | for d in [gtdir, generatedir]: 49 | os.makedirs(d, exist_ok=True) 50 | test_dataset = ImageDataset(opts) 51 | test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False, drop_last=False) 52 | state_path = os.path.join(ckpt_name) 53 | state_dict = torch.load(state_path) 54 | gen.load_state_dict(state_dict['gen']) 55 | gen.eval() 56 | with torch.no_grad(): 57 | for batch_idx, test_data in enumerate(test_loader): 58 | for k in test_data.keys(): 59 | if isinstance(test_data[k], torch.Tensor): 60 | test_data[k] = test_data[k].cuda() 61 | name = test_data['name'] 62 | fake = gen(test_data) 63 | gt = test_data['ground_truth'] 64 | for i in range(fake.size(0)): 65 | plt.imsave(os.path.join(gtdir, f'{name[i]}.png'), denorm_img(gt[i:i + 1]), vmin=0, vmax=1) 66 | plt.imsave(os.path.join(generatedir, f'{name[i]}.png'), denorm_img(fake[i:i + 1]), vmin=0, vmax=1) 67 | 68 | # FID socre https://github.com/mseitzer/pytorch-fid 69 | # inception score https://github.com/sbarratt/inception-score-pytorch 70 | os.system(f'python -m pytorch_fid {gtdir} {generatedir}') 71 | imgs = [] 72 | for f in os.listdir(generatedir): 73 | im = np.array(Image.open(os.path.join(generatedir, f))).transpose(2, 0, 1).astype(np.float32)[:3] 74 | im /= 255 75 | im = im * 2 - 1 76 | imgs.append(im) 77 | imgs = np.stack(imgs, 0) 78 | imgs = torch.from_numpy(imgs).cuda() 79 | iscore = inception_score(imgs, cuda=True, batch_size=32, resize=True, splits=1)[0] 80 | print('IS score', iscore) 81 | -------------------------------------------------------------------------------- /models/VITGen.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from .ops import get_sinusoid_encoding_table, CorssAttnBlock 6 | from .VIT import * 7 | from .QEM import QueryExpansionModule 8 | from .PSM import PatchSmoothingModule 9 | 10 | class TransGen(nn.Module): 11 | def __init__(self,opts, enc_ckpt_path=None): 12 | super(TransGen, self).__init__() 13 | self.output_size=opts.output_size 14 | self.input_size=opts.input_size 15 | self.patch_size=16 # same as ViT-B 16 | hidden_num=768 # same as ViT-B 17 | 18 | # initialize the weight of decoder, psm and qem 19 | self.qem=QueryExpansionModule(hidden_num=hidden_num,input_size=self.input_size,outout_size=self.output_size,patch_size=self.patch_size) 20 | self.transformer_decoder=nn.ModuleList([ 21 | CorssAttnBlock( 22 | dim=hidden_num, num_heads=12, mlp_ratio=4, qkv_bias=True, qk_scale=None, 23 | drop=0., attn_drop=0., drop_path=0., norm_layer=nn.LayerNorm, 24 | init_values=0., window_size= None) 25 | for _ in range(opts.dec_depth)]) 26 | self.psm=PatchSmoothingModule(patch_size=16,out_chans=3,embed_dim=hidden_num) 27 | self.apply(self._init_weights) 28 | 29 | # initialize the weight of encoder using pretrain checkpoint 30 | 31 | self.transformer_encoder = vit_base_patch16(pretrained=True, img_size=224, init_ckpt=enc_ckpt_path) 32 | #vit_base_patch16(pretrain=True, init_ckpt=enc_ckpt_path, img_size=self.input_size) 33 | 34 | self.enc_image_size=224 35 | 36 | # initialize the weight of encoder using pretrain checkpoint 37 | self.pos_embed = get_sinusoid_encoding_table(12**2, hidden_num) 38 | self.inner_index, self.outer_index=self.get_index() 39 | 40 | def get_index(self): 41 | input_query_width=self.input_size//self.patch_size 42 | output_query_width=self.output_size//self.patch_size 43 | mask=torch.ones(size=[output_query_width,output_query_width]).long() 44 | pad_width=(output_query_width-input_query_width)//2 45 | mask[pad_width:-pad_width,pad_width:-pad_width] = 0 46 | mask=mask.view(-1) 47 | return mask==0,mask==1 48 | 49 | def _init_weights(self, m): 50 | if isinstance(m, nn.Linear): 51 | nn.init.xavier_uniform_(m.weight) 52 | if isinstance(m, nn.Linear) and m.bias is not None: 53 | nn.init.constant_(m.bias, 0) 54 | elif isinstance(m, nn.LayerNorm): 55 | nn.init.constant_(m.bias, 0) 56 | nn.init.constant_(m.weight, 1.0) 57 | 58 | def forward(self, samples): 59 | if type(samples) is not dict: 60 | samples={'input':samples, 'gt_inner':F.pad(samples,(32,32,32,32))} 61 | x = samples['input'] 62 | 63 | gt_inner = samples['gt_inner'] 64 | 65 | b,c,w,h=x.size() 66 | 67 | assert w==128 and h==128 68 | padded_x = F.pad(x, (48, 48, 48, 48), mode='reflect') 69 | vit_mask = torch.ones(size=(14, 14)).long() 70 | vit_mask[3:-3, 3:-3] = 0 71 | 72 | vit_mask = vit_mask.view(-1).expand(b, -1).contiguous().bool() 73 | 74 | src = self.transformer_encoder.forward_features(padded_x, vit_mask) # b n c 75 | 76 | query_embed=self.qem(src) 77 | 78 | full_pos=self.pos_embed.type_as(x).to(x.device).clone().detach().expand(x.size(0),-1,-1) 79 | 80 | tgt_outer=query_embed[:,self.outer_index,:]+full_pos[:,self.outer_index,:] 81 | 82 | for i,dec in enumerate(self.transformer_decoder): 83 | tgt_outer = dec(tgt_outer, src) 84 | 85 | tgt = torch.zeros_like(query_embed,dtype=torch.float32) 86 | 87 | tgt[:, self.outer_index] = tgt_outer 88 | 89 | fake=self.psm(tgt,gt_inner) 90 | return fake 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # QueryOTR 2 | 3 | ## Outpainting by Queries, ECCV 2022. [Paper](https://link.springer.com/chapter/10.1007/978-3-031-20050-2_10) [ArXiv](https://arxiv.org/abs/2207.05312) 4 | 5 | we propose a novel hybrid vision-transformer-based encoder-decoder framework, named Query Outpainting TRansformer (QueryOTR), for extrapolating visual context all-side around a given image. Patch-wise mode's global modeling capacity allows us to extrapolate images from the attention mechanism's query standpoint. A novel Query Expansion Module (QEM) is designed to integrate information from the predicted queries based on the encoder's output, hence accelerating the convergence of the pure transformer even with a relatively small dataset. To further enhance connectivity between each patch, the proposed Patch Smoothing Module (PSM) re-allocates and averages the overlapped regions, thus providing seamless predicted images. We experimentally show that QueryOTR could generate visually appealing results smoothly and realistically against the state-of-the-art image outpainting approaches. 6 | 7 |
8 | 9 |
10 | 11 | ## 1. Requirements 12 | PyTorch >= 1.10.1; 13 | python >= 3.7; 14 | CUDA >= 11.3; 15 | torchvision; 16 | 17 | NOTE: The code was tested to work well on Linux with torch 1.7, 1.9 and Win10 with torch 1.10.1. However, there is potential "Inplace Operation Error" bug if you use PyTorch < 1.10, which is quiet subtle and not fixed. If you found why the bug occur, pls let me know. 18 | 19 | ## News: 20 | \[2022/11/7\] We update the code. We found the [official MAE](https://github.com/facebookresearch/mae) code may degrade the performance by unkonwn reason (about 0.5-1 in terms of FID), and we go back to [unofficial MAE](https://github.com/pengzhiliang/MAE-pytorch). Meanwhile, we upload a trained checkpoints on Scenery [google drive](https://drive.google.com/drive/folders/1s_Qs6m314a5vwLzdQ58uKOveK6fZjgaB?usp=share_link) which can reach FID 20.38, IS 3.959. It worth noting that the result may change due to the randomness of the code, e.g., one of the input of QEM is noise. 21 | 22 | 23 | ## 2. Data preparation 24 | 25 | ### Scenery 26 | Scenery consists of about 6,000 images, and we randomly select 1,000 images for evaluation. The training and test dataset can be down [here](https://github.com/z-x-yang/NS-Outpainting) 27 | 28 | Meanwhile, we also provide the Scenery dataset that we have split here [baidu_pan](https://pan.baidu.com/s/1Zn5X3jfqr6x3ho705VMHZA?pwd=qotr). 29 | 30 | ### Building 31 | Building contains about 16,000 images in the training set and 1,500 images in the testing set, which can be found in [here](https://github.com/PengleiGao/UTransformer) 32 | 33 | ### WikiArt 34 | The WikiArt datasets can be downloaded [here](https://github.com/cs-chan/ArtGAN/tree/master/WikiArt%20Dataset). We perform a split manner of genres datasets, which contains 45,503 training images and 19,492 testing images 35 | 36 | ## 3. Training and evaluation 37 | Before you reimplement our results, you need to download the ViT pretrain checkpoint [here](https://drive.google.com/drive/folders/1ZVzOD-ZGPBNtJ4HtsR-8IIH7Cm40LiMW?usp=share_link), and then initialize the encoder weight. 38 | 39 | 40 | Training on your datasets, run: 41 | ``` 42 | CUDA_VISIBLE_DEVICES= python main.py --name=EXPERIMENT_NAME --data_root=YOUR_TRAIN_PATH --patch_mean=YOUR_PATCH_MEAN --patch_std=YOUR_PATCH_STD 43 | ``` 44 | 45 | Evaluate on your datasets, run: 46 | ``` 47 | CUDA_VISIBLE_DEVICES= python evaluate.py --r=EXPERIMENT_NAME --data_root=YOUR_TEST_PATH --patch_mean=YOUR_PATCH_MEAN --patch_std=YOUR_PATCH_STD 48 | ``` 49 | 50 | 51 | 52 | 53 | ## Acknowledgements 54 | 55 | Our codes are built upon MAE, [pytroch-fid](https://github.com/mseitzer/pytorch-fid) and [inception score](https://github.com/sbarratt/inception-score-pytorch) 56 | 57 | ## Citation 58 | 59 | ``` 60 | @inproceedings{yao2022qotr, 61 | title={Outpainting by Queries}, 62 | author={Yao, Kai and Gao, Penglei and Yang, Xi and Sun, Jie and Zhang, Rui and Huang, Kaizhu}, 63 | booktitle={European Conference on Computer Vision}, 64 | pages={153--169}, 65 | year={2022}, 66 | organization={Springer} 67 | } 68 | ``` 69 | -------------------------------------------------------------------------------- /models/QEM.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torchvision.ops import DeformConv2d, deform_conv2d 4 | 5 | 6 | class DeformConv(DeformConv2d): 7 | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, bias=None): 8 | super(DeformConv, self).__init__(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, 9 | stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias) 10 | channels_ = groups * 3 * self.kernel_size[0] * self.kernel_size[1] 11 | self.conv_offset_mask = nn.Conv2d(self.in_channels, 12 | channels_, 13 | kernel_size=self.kernel_size, 14 | stride=self.stride, 15 | padding=self.padding, 16 | bias=True) 17 | self.init_offset() 18 | 19 | def init_offset(self): 20 | self.conv_offset_mask.weight.data.zero_() 21 | self.conv_offset_mask.bias.data.zero_() 22 | 23 | def forward(self, input): 24 | out = self.conv_offset_mask(input) 25 | o1, o2, mask = torch.chunk(out, 3, dim=1) 26 | offset = torch.cat((o1, o2), dim=1) 27 | mask = torch.sigmoid(mask) 28 | return deform_conv2d(input, offset, self.weight, self.bias, stride=self.stride, 29 | padding=self.padding, dilation=self.dilation, mask=mask) 30 | 31 | 32 | class ResidualBlock(nn.Module): 33 | def __init__(self, planes, ): 34 | super(ResidualBlock, self).__init__() 35 | self.conv1 = nn.Conv2d(planes, planes, 3, 1, 1) 36 | self.conv2 = DeformConv(planes, planes, 3, 1, 1) 37 | self.norm1 = nn.InstanceNorm2d(planes, affine=True) 38 | self.norm2 = nn.InstanceNorm2d(planes, affine=True) 39 | self.act = nn.LeakyReLU(0.2, True) 40 | 41 | def forward(self, x): 42 | x_sc = x 43 | x = self.norm1(x) 44 | x = self.act(x) 45 | x = self.conv1(x) 46 | x = self.norm2(x) 47 | x = self.act(x) 48 | x = self.conv2(x) 49 | return x + x_sc 50 | 51 | 52 | class QueryExpansionModule(nn.Module): 53 | def __init__(self, hidden_num=768, n_block=8, input_size=128, outout_size=192, patch_size=16): 54 | super(QueryExpansionModule, self).__init__() 55 | 56 | self.hidden_num = hidden_num 57 | self.input_query_width = input_size // patch_size 58 | self.output_query_width = outout_size // patch_size 59 | self.res_blocks = nn.ModuleList([ResidualBlock(hidden_num) for _ in range(n_block)]) 60 | self.noise_mlp = nn.Sequential(nn.Linear(hidden_num // 8, hidden_num // 4), nn.LayerNorm(hidden_num // 4), 61 | nn.ReLU(), 62 | nn.Linear(hidden_num // 4, hidden_num // 2), nn.LayerNorm(hidden_num // 2), 63 | nn.ReLU(), 64 | nn.Linear(hidden_num // 2, hidden_num)) 65 | 66 | self.norm = nn.LayerNorm(hidden_num) 67 | self.embed = nn.Linear(hidden_num, hidden_num) 68 | self.inner_query_index, self.outer_query_index = self.get_index() 69 | 70 | def get_index(self): 71 | mask = torch.ones(size=[self.output_query_width, self.output_query_width]).long() 72 | pad_width = (self.output_query_width - self.input_query_width) // 2 73 | mask[pad_width:-pad_width, pad_width:-pad_width] = 0 74 | mask = mask.view(-1) 75 | return mask == 0, mask == 1 76 | 77 | def forward(self, src_query): 78 | b, n, c = src_query.size() 79 | 80 | ori_src_query = src_query 81 | assert src_query.size( 82 | 1) == self.input_query_width ** 2, f'QEM input spatial dimension is wrong, {src_query.size(1)} and {self.input_query_width ** 2}' 83 | noise = torch.randn(size=(b, self.output_query_width ** 2, c // 8), dtype=torch.float32).to(src_query.device) 84 | 85 | initial_query = self.noise_mlp(noise) 86 | 87 | initial_query[:, self.inner_query_index] = src_query 88 | 89 | x = initial_query.permute(0, 2, 1).reshape(b, c, self.output_query_width, self.output_query_width).contiguous() 90 | 91 | for layer in self.res_blocks: 92 | x = layer(x) 93 | 94 | x = x.flatten(2).transpose(1, 2) 95 | x = self.norm(x) 96 | x[:, self.inner_query_index, :] = ori_src_query 97 | x = self.embed(x) 98 | return x 99 | 100 | if __name__ == '__main__': 101 | m1 = QueryExpansionModule() 102 | x1 = torch.randn([1, 64, 768]) 103 | y1 = m1(x1) 104 | print(y1.size()) 105 | -------------------------------------------------------------------------------- /models/VIT_old.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | # -------------------------------------------------------- 7 | # References: 8 | # timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm 9 | # DeiT: https://github.com/facebookresearch/deit 10 | # -------------------------------------------------------- 11 | 12 | from functools import partial 13 | 14 | import torch 15 | import torch.nn as nn 16 | 17 | import timm.models.vision_transformer 18 | 19 | 20 | class VisionTransformer(timm.models.vision_transformer.VisionTransformer): 21 | """ Vision Transformer with support for global average pooling 22 | """ 23 | def __init__(self, global_pool=False, **kwargs): 24 | super(VisionTransformer, self).__init__(**kwargs) 25 | 26 | self.global_pool = global_pool 27 | if self.global_pool: 28 | norm_layer = kwargs['norm_layer'] 29 | embed_dim = kwargs['embed_dim'] 30 | self.fc_norm = norm_layer(embed_dim) 31 | 32 | del self.norm # remove the original norm 33 | 34 | 35 | def forward_features(self, x): 36 | x = self.patch_embed(x) 37 | 38 | #cls_tokens = self.cls_token.expand(B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks 39 | #x = torch.cat((cls_tokens, x), dim=1) 40 | # exclude the cls token and its corresponding position embeding 41 | x = x + self.pos_embed[:, 1:] 42 | x = self.pos_drop(x) 43 | 44 | for blk in self.blocks: 45 | x = blk(x) 46 | return x 47 | 48 | def interpolate_pos_embed(model, checkpoint_model): 49 | if 'pos_embed' in checkpoint_model: 50 | pos_embed_checkpoint = checkpoint_model['pos_embed'] 51 | embedding_size = pos_embed_checkpoint.shape[-1] 52 | num_patches = model.patch_embed.num_patches 53 | num_extra_tokens = model.pos_embed.shape[-2] - num_patches 54 | # height (== width) for the checkpoint position embedding 55 | orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5) 56 | # height (== width) for the new position embedding 57 | new_size = int(num_patches ** 0.5) 58 | # class_token and dist_token are kept unchanged 59 | if orig_size != new_size: 60 | print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size)) 61 | extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens] 62 | # only the position tokens are interpolated 63 | pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:] 64 | pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2) 65 | pos_tokens = torch.nn.functional.interpolate( 66 | pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False) 67 | pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2) 68 | new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1) 69 | checkpoint_model['pos_embed'] = new_pos_embed 70 | 71 | from timm.models.layers import trunc_normal_ 72 | def vit_base_patch16(pretrain=False, init_ckpt=None, **kwargs): 73 | model = VisionTransformer( 74 | patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, 75 | norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) 76 | if pretrain and init_ckpt is not None: 77 | checkpoint = torch.load(init_ckpt, map_location='cpu') 78 | 79 | print("Load pre-trained checkpoint from: %s" % init_ckpt) 80 | checkpoint_model = checkpoint['model'] 81 | state_dict = model.state_dict() 82 | for k in ['head.weight', 'head.bias']: 83 | if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape: 84 | print(f"Removing key {k} from pretrained checkpoint") 85 | del checkpoint_model[k] 86 | 87 | # interpolate position embedding 88 | interpolate_pos_embed(model, checkpoint_model) 89 | 90 | # load pre-trained model 91 | msg = model.load_state_dict(checkpoint_model, strict=False) 92 | print(msg) 93 | 94 | # manually initialize fc layer 95 | trunc_normal_(model.head.weight, std=2e-5) 96 | return model 97 | 98 | 99 | def vit_large_patch16(**kwargs): 100 | model = VisionTransformer( 101 | patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, 102 | norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) 103 | return model 104 | 105 | 106 | def vit_huge_patch14(**kwargs): 107 | model = VisionTransformer( 108 | patch_size=14, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True, 109 | norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) 110 | return model 111 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' 4 | import torch 5 | import itertools 6 | import datetime 7 | 8 | torch.backends.cudnn.benchmark = True 9 | from torch.utils.data import DataLoader 10 | from datasets import ImageDataset 11 | from util.misc import cosine_scheduler 12 | 13 | import argparse 14 | 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument('--name', type=str, default='scenery') 17 | 18 | parser.add_argument('--lr', type=float, default=1e-4) 19 | parser.add_argument('--min_lr', type=float, default=1e-4) 20 | parser.add_argument('--warnup_epoch', type=int, default=10) 21 | parser.add_argument('--max_epoch', type=int, default=300) 22 | parser.add_argument('--batch_size', type=int, default=64) 23 | parser.add_argument('--num_workers', type=int, default=8) 24 | 25 | parser.add_argument('--eval', default=False, type=bool) 26 | parser.add_argument('--half_precision', default=False, type=bool) 27 | 28 | parser.add_argument('--input_size', type=int, default=128) 29 | parser.add_argument('--output_size', type=int, default=192) 30 | parser.add_argument('--enc_ckpt_path', type=str, default='pretrain_mae_vit_base_mask_0.75_400e.pth') 31 | parser.add_argument('--dec_depth', type=int, default=4) 32 | 33 | parser.add_argument('--data_root', type=str, default='E:/data3/train') 34 | parser.add_argument('--normlize_target', default=True, type=bool, help='normalized the target patch pixels') 35 | parser.add_argument('--patch_mean', type=float, default=0.5044838) 36 | parser.add_argument('--patch_std', type=float, default=0.1355051) 37 | 38 | from models.VITGen import TransGen 39 | from models.CNNDis import MsImageDis 40 | from losses import SetCriterion 41 | from engine import train_one_epoch, train_one_epoch_warmup 42 | 43 | if __name__ == '__main__': 44 | opts = parser.parse_args() 45 | 46 | train_dataset = ImageDataset(opts) 47 | train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=opts.batch_size, 48 | num_workers=opts.num_workers, persistent_workers=opts.num_workers > 0, 49 | shuffle=True, pin_memory=True) 50 | 51 | gen = TransGen(opts=opts, enc_ckpt_path=opts.enc_ckpt_path).cuda() 52 | cnn_dis = MsImageDis().cuda() 53 | 54 | g_param_dicts = [ 55 | {"params": [p for n, p in gen.named_parameters() if 'conv_offset_mask' not in n and not 'transformer_encoder' in n], "lr_scale": 1}, 56 | {"params": [p for n, p in gen.named_parameters() if 'conv_offset_mask' in n], "lr_scale": 0.1}, 57 | {"params": [p for n, p in gen.named_parameters() if 'transformer_encoder' in n], "lr_scale": 1} 58 | ] 59 | 60 | opt_g = torch.optim.Adam(g_param_dicts, lr=opts.lr, betas=(0.0, 0.99), weight_decay=1e-4) 61 | opt_d = torch.optim.Adam(itertools.chain(cnn_dis.parameters()), lr=opts.lr, betas=(0.0, 0.99), weight_decay=1e-4) 62 | lr_schedule_values = cosine_scheduler(opts.lr, opts.min_lr, opts.max_epoch, len(train_loader), 63 | warmup_epochs=opts.warnup_epoch, warmup_steps=-1) 64 | 65 | now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S") 66 | nowname = now + '_' + opts.name 67 | logdir = os.path.join("logs", nowname) 68 | ckptdir = os.path.join(logdir, "checkpoints") 69 | visdir = os.path.join(logdir, "visuals") 70 | for d in [logdir, ckptdir, visdir]: 71 | os.makedirs(d, exist_ok=True) 72 | opts.visdir = visdir 73 | opts.ckptdir = ckptdir 74 | 75 | if opts.half_precision: 76 | g_grad_scaler = torch.cuda.amp.GradScaler() 77 | else: 78 | g_grad_scaler = None 79 | 80 | criterion = SetCriterion(opts) 81 | 82 | iteration = 1 83 | for epoch in range(opts.max_epoch): 84 | # warm up the learning rate 85 | if lr_schedule_values is not None and epoch < opts.warnup_epoch: 86 | for i, param_group in enumerate(opt_g.param_groups): 87 | param_group["lr"] = lr_schedule_values[iteration] * param_group["lr_scale"] 88 | for i, param_group in enumerate(opt_d.param_groups): 89 | param_group["lr"] = lr_schedule_values[iteration] 90 | else: 91 | for i, param_group in enumerate(opt_g.param_groups): 92 | param_group["lr"] = opts.lr * param_group["lr_scale"] 93 | for i, param_group in enumerate(opt_d.param_groups): 94 | param_group["lr"] = opts.lr 95 | 96 | if epoch < opts.warnup_epoch: 97 | train_one_epoch_warmup(opts, gen, criterion, train_loader, opt_g, torch.device('cuda'), epoch, 98 | g_grad_scale=g_grad_scaler) 99 | else: 100 | train_one_epoch(opts, gen, cnn_dis, criterion, train_loader, opt_g, opt_d, torch.device('cuda'), epoch, 101 | g_grad_scale=g_grad_scaler) 102 | 103 | iteration += len(train_loader) 104 | 105 | if (epoch + 1) % 10 == 0 and epoch > 50: 106 | torch.save({'gen': gen.state_dict()}, os.path.join(ckptdir, f'{epoch + 1}.pth')) 107 | torch.save({'gen': gen.state_dict()}, os.path.join(ckptdir, f'latest.pth')) 108 | -------------------------------------------------------------------------------- /models/PSM.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch 3 | from einops import rearrange 4 | import torch.nn.functional as F 5 | 6 | # Equivalent implementation in the paper for efficiency 7 | class PatchSmoothingModule(nn.Module): 8 | def __init__(self, embed_dim=768, out_chans=3, input_size=128, output_size=192, patch_size=16, overlap_size=8, 9 | bias=True): 10 | super().__init__() 11 | self.use_bias = bias 12 | self.patch_size = patch_size 13 | self.input_size = input_size 14 | self.output_size = output_size 15 | 16 | self.embed_dim = embed_dim 17 | patch_size = patch_size 18 | kernel_size = patch_size + overlap_size * 2 19 | padding_size = overlap_size 20 | self.proj = nn.ConvTranspose2d(embed_dim, out_chans, bias=False, kernel_size=kernel_size, stride=patch_size, 21 | padding=padding_size) 22 | if bias: 23 | self.bias = torch.nn.Parameter(torch.FloatTensor(1, out_chans, kernel_size, kernel_size), 24 | requires_grad=True) 25 | nn.init.constant_(self.bias, 0) 26 | 27 | self.mask = torch.ones(1, 1, output_size // patch_size, output_size // patch_size) 28 | p = ((output_size - input_size) // 2) // patch_size 29 | self.mask[:, :, p:-p, p:-p] = 0 30 | 31 | self.mask_weight = F.conv_transpose2d(self.mask.detach(), torch.ones([1, out_chans, kernel_size, kernel_size]), 32 | bias=None, stride=patch_size, padding=padding_size) 33 | self.mask_weight[self.mask_weight != 0] = 1 / self.mask_weight[self.mask_weight != 0] 34 | self.patch_size = patch_size 35 | self.padding_size = padding_size 36 | 37 | def forward(self, x, gt_inner): 38 | assert x.size(1) == (self.output_size // self.patch_size) ** 2 39 | x = rearrange(x, 'b (h w) c -> b c h w', h=self.output_size // self.patch_size) 40 | x = self.proj(x) 41 | 42 | if self.use_bias: 43 | bias = F.conv_transpose2d(self.mask.detach().to(x.device), self.bias, bias=None, stride=self.patch_size, 44 | padding=self.padding_size) 45 | x = x + bias 46 | x = x * self.mask_weight.to(x.device) 47 | p = (self.output_size - self.input_size) // 2 48 | x[:, :, p:-p, p:-p] = gt_inner[:, :, p:-p, p:-p] 49 | 50 | return x 51 | 52 | # Original implementation in the paper 53 | class PatchSmoothingModule_v2(nn.Module): 54 | def __init__(self, embed_dim=768, out_chans=3, input_size=128, output_size=192, patch_size=16, overlap_size=8, 55 | bias=True): 56 | super().__init__() 57 | self.use_bias = bias 58 | self.patch_size = patch_size 59 | self.input_size = input_size 60 | self.output_size = output_size 61 | self.out_chans = out_chans 62 | self.embed_dim = embed_dim 63 | self.overlap_size = overlap_size 64 | 65 | patch_size = patch_size 66 | self.kernel_size = kernel_size = patch_size + overlap_size * 2 67 | padding_size = overlap_size 68 | self.proj = nn.ConvTranspose2d(embed_dim, out_chans, bias=bias, kernel_size=kernel_size, stride=kernel_size, 69 | padding=0) 70 | 71 | self.patch_size = patch_size 72 | self.padding_size = padding_size 73 | 74 | def forward(self, x, gt_inner): 75 | assert x.size(1) == (self.output_size // self.patch_size) ** 2 76 | x = rearrange(x, 'b (h w) c -> (b h w) c', h=self.output_size // self.patch_size) 77 | x = x.unsqueeze(-1).unsqueeze(-1) 78 | x = self.proj(x) # (b h w) c 32 32 79 | x = rearrange(x, '(b h w) c p1 p2-> b c h w p1 p2', h=self.output_size // self.patch_size, 80 | w=self.output_size // self.patch_size) 81 | output_patches = x 82 | 83 | output = torch.zeros([x.size(0), self.out_chans, self.output_size + 2 * self.overlap_size, 84 | self.output_size + 2 * self.overlap_size]) 85 | 86 | mask_weight = torch.zeros( 87 | [1, 1, self.output_size + 2 * self.overlap_size, self.output_size + 2 * self.overlap_size]) 88 | for i in range(self.output_size // self.patch_size): 89 | for j in range(self.output_size // self.patch_size): 90 | output[:, :, i * self.patch_size: (i + 1) * self.patch_size + self.overlap_size * 2, 91 | j * self.patch_size: (j + 1) * self.patch_size + self.overlap_size * 2] += x[:, :, i, j, :, :] 92 | mask_weight[:, :, i * self.patch_size: (i + 1) * self.patch_size + self.overlap_size * 2, 93 | j * self.patch_size: (j + 1) * self.patch_size + self.overlap_size * 2] += 1 94 | 95 | mask_weight[mask_weight != 0] = 1 / mask_weight[mask_weight != 0] 96 | output = output * mask_weight 97 | output = output[:, :, self.overlap_size:-self.overlap_size, self.overlap_size:-self.overlap_size] 98 | p = (self.output_size - self.input_size) // 2 99 | output[:, :, p:-p, p:-p] = gt_inner[:, :, p:-p, p:-p] 100 | return output, output_patches 101 | 102 | if __name__ == '__main__': 103 | m1 = PatchSmoothingModule() 104 | m2 = PatchSmoothingModule_v2() 105 | x1 = torch.randn([1, 144, 768]) 106 | x2 = torch.randn([1, 3, 192, 192]) 107 | y1 = m1(x1, x2) 108 | y2, _ = m2(x1, x2) 109 | print(y1.size(), y2.size()) 110 | 111 | -------------------------------------------------------------------------------- /engine.py: -------------------------------------------------------------------------------- 1 | 2 | import math 3 | import os 4 | import sys 5 | from typing import Iterable 6 | import torch 7 | 8 | import functools 9 | print = functools.partial(print, flush=True) 10 | import util.misc as utils 11 | 12 | from einops import rearrange 13 | import numpy as np 14 | import matplotlib.pyplot as plt 15 | import time 16 | 17 | def denorm_img(tensor, opts): 18 | tensor = rearrange(tensor[0:4], 'b c w h -> b w h c').detach().cpu() 19 | tensor = tensor * torch.tensor((opts.patch_std,opts.patch_std,opts.patch_std)) + torch.tensor((opts.patch_mean,opts.patch_mean,opts.patch_mean)) 20 | tensor = np.clip(tensor.flatten(0, 1).numpy(), 0, 1) 21 | return tensor 22 | 23 | def train_one_epoch(opts, GEN: torch.nn.Module, DIS: torch.nn.Module, criterion: torch.nn.Module, 24 | data_loader: Iterable, gen_opt: torch.optim.Optimizer, dis_opt: torch.optim.Optimizer, 25 | device: torch.device, epoch: int, g_grad_scale=None): 26 | GEN.train() 27 | DIS.train() 28 | criterion.train() 29 | metric_logger = utils.MetricLogger(delimiter=" ") 30 | metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}')) 31 | 32 | header = 'Epoch: [{}]'.format(epoch) 33 | print_freq = 10 34 | display_freq=75 35 | 36 | for i, samples in enumerate(metric_logger.log_every(data_loader, print_freq, header)): 37 | for k,v in samples.items(): 38 | if isinstance(samples[k], torch.Tensor): 39 | samples[k]=v.to(device) 40 | if g_grad_scale is None: 41 | gt=samples['ground_truth'] 42 | fake = GEN(samples) 43 | if i%display_freq==0: 44 | fig_name = f"{epoch}_{time.time():04f}" 45 | fig = np.concatenate([denorm_img(gt, opts), denorm_img(fake, opts)], axis=1) 46 | plt.imsave(os.path.join(opts.visdir, fig_name + '.png'), fig, vmin=0, vmax=1) 47 | 48 | D_loss_dict = criterion.get_dis_loss(fake, gt, DIS) 49 | D_losses = sum(D_loss_dict[k] * criterion.dis_weight_dict[k] for k in D_loss_dict.keys()) 50 | 51 | dis_opt.zero_grad() 52 | D_losses.backward() 53 | dis_opt.step() 54 | 55 | G_loss_dict = criterion.get_gen_loss(fake, gt, DIS) 56 | G_losses = sum(G_loss_dict[k] * criterion.gen_weight_dict[k] for k in G_loss_dict.keys()) 57 | 58 | gen_opt.zero_grad() 59 | G_losses.backward() 60 | gen_opt.step() 61 | 62 | else: 63 | with torch.cuda.amp.autocast(): 64 | gt = samples['ground_truth'] 65 | fake = GEN(samples) 66 | 67 | if i % display_freq == 0: 68 | fig_name = f"{epoch}_{time.time():04f}" 69 | fig = np.concatenate([denorm_img(gt, opts), denorm_img(fake, opts)], axis=1) 70 | plt.imsave(os.path.join(opts.visdir, fig_name + '.png'), fig, vmin=0, vmax=1) 71 | 72 | D_loss_dict = criterion.get_dis_loss(fake, gt) 73 | D_losses = sum(D_loss_dict[k] * criterion.dis_weight_dict[k] for k in D_loss_dict.keys()) 74 | 75 | dis_opt.zero_grad() 76 | D_losses.backward() 77 | dis_opt.step() 78 | 79 | with torch.cuda.amp.autocast(): 80 | G_loss_dict = criterion.get_gen_loss(fake, gt) 81 | G_losses = sum(G_loss_dict[k] * criterion.gen_weight_dict[k] for k in G_loss_dict.keys()) 82 | 83 | gen_opt.zero_grad() 84 | g_grad_scale.scale(G_losses).backward() 85 | g_grad_scale.step(gen_opt) 86 | g_grad_scale.update() 87 | 88 | metric_logger.update(**G_loss_dict,**D_loss_dict) 89 | metric_logger.update(lr=dis_opt.param_groups[0]["lr"]) 90 | 91 | metric_logger.synchronize_between_processes() 92 | print("Averaged stats:", metric_logger) 93 | return {k: meter.global_avg for k, meter in metric_logger.meters.items()} 94 | 95 | 96 | def train_one_epoch_warmup(opts, GEN: torch.nn.Module,criterion: torch.nn.Module, 97 | data_loader: Iterable, gen_opt: torch.optim.Optimizer, 98 | device: torch.device, epoch: int, g_grad_scale=None): 99 | 100 | 101 | GEN.train() 102 | criterion.train() 103 | metric_logger = utils.MetricLogger(delimiter=" ") 104 | metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}')) 105 | 106 | header = '(Warning Up!!)Epoch: [{}]'.format(epoch) 107 | print_freq = 10 108 | display_freq = 75 109 | 110 | for i,samples in enumerate(metric_logger.log_every(data_loader, print_freq, header)): 111 | for k, v in samples.items(): 112 | if isinstance(samples[k], torch.Tensor): 113 | samples[k] = v.to(device) 114 | 115 | if g_grad_scale is None: 116 | gt = samples['ground_truth'] 117 | outputs = GEN(samples) 118 | if i%display_freq==0: 119 | fig_name = f"{epoch}_{time.time():04f}" 120 | fig = np.concatenate([denorm_img(gt, opts), denorm_img(outputs, opts)], axis=1) 121 | plt.imsave(os.path.join(opts.visdir, fig_name + '.png'), fig, vmin=0, vmax=1) 122 | G_loss_dict = criterion.get_gen_loss(outputs, gt, warmup=True) 123 | G_losses = sum(G_loss_dict[k] * criterion.gen_weight_dict[k] for k in G_loss_dict.keys()) 124 | 125 | gen_opt.zero_grad() 126 | G_losses.backward() 127 | gen_opt.step() 128 | else: 129 | gen_opt.zero_grad() 130 | with torch.cuda.amp.autocast(): 131 | gt = samples['ground_truth'] 132 | outputs = GEN(samples) 133 | if i % display_freq == 0: 134 | fig_name = f"{epoch}_{time.time():04f}" 135 | fig = np.concatenate([denorm_img(gt, opts), denorm_img(outputs, opts)], axis=1) 136 | plt.imsave(os.path.join(opts.visdir, fig_name + '.png'), fig, vmin=0, vmax=1) 137 | G_loss_dict = criterion.get_gen_loss(outputs, gt, warmup=True) 138 | G_losses = sum(G_loss_dict[k] * criterion.gen_weight_dict[k] for k in G_loss_dict.keys()) 139 | g_grad_scale.scale(G_losses).backward() 140 | g_grad_scale.step(gen_opt) 141 | g_grad_scale.update() 142 | 143 | metric_logger.update(**G_loss_dict) 144 | metric_logger.update(lr=gen_opt.param_groups[0]["lr"]) 145 | 146 | metric_logger.synchronize_between_processes() 147 | print("Averaged stats:", metric_logger) 148 | return {k: meter.global_avg for k, meter in metric_logger.meters.items()} -------------------------------------------------------------------------------- /util/misc.py: -------------------------------------------------------------------------------- 1 | import time 2 | from collections import defaultdict, deque 3 | import datetime 4 | import functools 5 | print = functools.partial(print, flush=True) 6 | 7 | import torch 8 | import torch.distributed as dist 9 | import numpy as np 10 | import math 11 | 12 | def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0, 13 | start_warmup_value=0, warmup_steps=-1): 14 | warmup_schedule = np.array([]) 15 | warmup_iters = warmup_epochs * niter_per_ep 16 | if warmup_steps > 0: 17 | warmup_iters = warmup_steps 18 | print("Set warmup steps = %d" % warmup_iters) 19 | if warmup_epochs > 0: 20 | warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters) 21 | 22 | iters = np.arange(epochs * niter_per_ep - warmup_iters) 23 | schedule = np.array( 24 | [final_value + 0.5 * (base_value - final_value) * (1 + math.cos(math.pi * i / (len(iters)))) for i in iters]) 25 | 26 | schedule = np.concatenate((warmup_schedule, schedule)) 27 | 28 | assert len(schedule) == epochs * niter_per_ep 29 | return schedule 30 | 31 | def is_dist_avail_and_initialized(): 32 | if not dist.is_available(): 33 | return False 34 | if not dist.is_initialized(): 35 | return False 36 | return True 37 | 38 | class SmoothedValue(object): 39 | """Track a series of values and provide access to smoothed values over a 40 | window or the global series average. 41 | """ 42 | 43 | def __init__(self, window_size=20, fmt=None): 44 | if fmt is None: 45 | fmt = "{median:.4f} ({global_avg:.4f})" 46 | self.deque = deque(maxlen=window_size) 47 | self.total = 0.0 48 | self.count = 0 49 | self.fmt = fmt 50 | 51 | def update(self, value, n=1): 52 | self.deque.append(value) 53 | self.count += n 54 | self.total += value * n 55 | 56 | def synchronize_between_processes(self): 57 | """ 58 | Warning: does not synchronize the deque! 59 | """ 60 | if not is_dist_avail_and_initialized(): 61 | return 62 | t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda') 63 | dist.barrier() 64 | dist.all_reduce(t) 65 | t = t.tolist() 66 | self.count = int(t[0]) 67 | self.total = t[1] 68 | 69 | @property 70 | def median(self): 71 | d = torch.tensor(list(self.deque)) 72 | return d.median().item() 73 | 74 | @property 75 | def avg(self): 76 | d = torch.tensor(list(self.deque), dtype=torch.float32) 77 | return d.mean().item() 78 | 79 | @property 80 | def global_avg(self): 81 | return self.total / self.count 82 | 83 | @property 84 | def max(self): 85 | return max(self.deque) 86 | 87 | @property 88 | def value(self): 89 | return self.deque[-1] 90 | 91 | def __str__(self): 92 | return self.fmt.format( 93 | median=self.median, 94 | avg=self.avg, 95 | global_avg=self.global_avg, 96 | max=self.max, 97 | value=self.value) 98 | 99 | class MetricLogger(object): 100 | def __init__(self, delimiter="\t"): 101 | self.meters = defaultdict(SmoothedValue) 102 | self.delimiter = delimiter 103 | 104 | def update(self, **kwargs): 105 | for k, v in kwargs.items(): 106 | if isinstance(v, torch.Tensor): 107 | v = v.item() 108 | assert isinstance(v, (float, int)) 109 | self.meters[k].update(v) 110 | 111 | def __getattr__(self, attr): 112 | if attr in self.meters: 113 | return self.meters[attr] 114 | if attr in self.__dict__: 115 | return self.__dict__[attr] 116 | raise AttributeError("'{}' object has no attribute '{}'".format( 117 | type(self).__name__, attr)) 118 | 119 | def __str__(self): 120 | loss_str = [] 121 | for name, meter in self.meters.items(): 122 | loss_str.append( 123 | "{}: {}".format(name, str(meter)) 124 | ) 125 | return self.delimiter.join(loss_str) 126 | 127 | def synchronize_between_processes(self): 128 | for meter in self.meters.values(): 129 | meter.synchronize_between_processes() 130 | 131 | def add_meter(self, name, meter): 132 | self.meters[name] = meter 133 | 134 | def log_every(self, iterable, print_freq, header=None): 135 | i = 0 136 | if not header: 137 | header = '' 138 | start_time = time.time() 139 | end = time.time() 140 | iter_time = SmoothedValue(fmt='{avg:.4f}') 141 | data_time = SmoothedValue(fmt='{avg:.4f}') 142 | space_fmt = ':' + str(len(str(len(iterable)))) + 'd' 143 | if torch.cuda.is_available(): 144 | log_msg = self.delimiter.join([ 145 | header, 146 | '[{0' + space_fmt + '}/{1}]', 147 | 'eta: {eta}', 148 | '{meters}', 149 | 'time: {time}', 150 | 'data: {data}', 151 | 'max mem: {memory:.0f}' 152 | ]) 153 | else: 154 | log_msg = self.delimiter.join([ 155 | header, 156 | '[{0' + space_fmt + '}/{1}]', 157 | 'eta: {eta}', 158 | '{meters}', 159 | 'time: {time}', 160 | 'data: {data}' 161 | ]) 162 | MB = 1024.0 * 1024.0 163 | for obj in iterable: 164 | data_time.update(time.time() - end) 165 | yield obj 166 | iter_time.update(time.time() - end) 167 | if i % print_freq == 0 or i == len(iterable) - 1: 168 | eta_seconds = iter_time.global_avg * (len(iterable) - i) 169 | eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) 170 | if is_main_process(): 171 | if torch.cuda.is_available(): 172 | print(log_msg.format( 173 | i, len(iterable), eta=eta_string, 174 | meters=str(self), 175 | time=str(iter_time), data=str(data_time), 176 | memory=torch.cuda.max_memory_allocated() / MB)) 177 | else: 178 | print(log_msg.format( 179 | i, len(iterable), eta=eta_string, 180 | meters=str(self), 181 | time=str(iter_time), data=str(data_time))) 182 | i += 1 183 | end = time.time() 184 | total_time = time.time() - start_time 185 | total_time_str = str(datetime.timedelta(seconds=int(total_time))) 186 | if is_main_process(): 187 | print('{} Total time: {} ({:.4f} s / it)'.format( 188 | header, total_time_str, total_time / len(iterable))) 189 | 190 | def is_main_process(): 191 | return get_rank() == 0 192 | 193 | def get_rank(): 194 | if not is_dist_avail_and_initialized(): 195 | return 0 196 | return dist.get_rank() 197 | 198 | -------------------------------------------------------------------------------- /models/CNNDis.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | from torch import nn 6 | from torch.autograd import Variable 7 | import torch 8 | import torch.nn.functional as F 9 | try: 10 | from itertools import izip as zip 11 | except ImportError: # will be 3.x series 12 | pass 13 | from torch.nn.utils.parametrizations import spectral_norm 14 | ################################################################################## 15 | # Discriminator 16 | ################################################################################## 17 | class Conv2dBlock(nn.Module): 18 | def __init__(self, input_dim ,output_dim, kernel_size, stride, 19 | padding=0, norm='none', activation='relu', pad_type='zero'): 20 | super(Conv2dBlock, self).__init__() 21 | self.use_bias = True 22 | # initialize padding 23 | if pad_type == 'reflect': 24 | self.pad = nn.ReflectionPad2d(padding) 25 | elif pad_type == 'replicate': 26 | self.pad = nn.ReplicationPad2d(padding) 27 | elif pad_type == 'zero': 28 | self.pad = nn.ZeroPad2d(padding) 29 | else: 30 | assert 0, "Unsupported padding type: {}".format(pad_type) 31 | 32 | # initialize convolution 33 | if 'sn' in norm: 34 | self.conv = spectral_norm(nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias)) 35 | norm=norm.replace('sn','') 36 | else: 37 | self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias) 38 | 39 | # initialize normalization 40 | norm_dim = output_dim 41 | if norm == 'bn': 42 | self.norm = nn.BatchNorm2d(norm_dim) 43 | elif norm == 'in': 44 | self.norm = nn.InstanceNorm2d(norm_dim) 45 | elif norm == 'none': 46 | self.norm = None 47 | else: 48 | assert 0, "Unsupported normalization: {}".format(norm) 49 | 50 | # initialize activation 51 | if activation == 'relu': 52 | self.activation = nn.ReLU(inplace=True) 53 | elif activation == 'lrelu': 54 | self.activation = nn.LeakyReLU(0.2, inplace=True) 55 | elif activation == 'prelu': 56 | self.activation = nn.PReLU() 57 | elif activation == 'selu': 58 | self.activation = nn.SELU(inplace=True) 59 | elif activation == 'tanh': 60 | self.activation = nn.Tanh() 61 | elif activation == 'none': 62 | self.activation = None 63 | else: 64 | assert 0, "Unsupported activation: {}".format(activation) 65 | 66 | def forward(self, x): 67 | x = self.conv(self.pad(x)) 68 | if self.norm: 69 | x = self.norm(x) 70 | if self.activation: 71 | x = self.activation(x) 72 | return x 73 | 74 | from .DiffAug import DiffAugment 75 | 76 | class MsImageDis(nn.Module): 77 | # Multi-scale discriminator architecture 78 | def __init__(self,input_dim=3,n_layer=4,num_scales=2): 79 | super(MsImageDis, self).__init__() 80 | self.n_layer = n_layer 81 | self.gan_type = 'hinge' 82 | self.dim = 64 83 | self.norm = 'snin' 84 | self.activ = 'lrelu' 85 | self.num_scales = num_scales 86 | self.pad_type = 'reflect' 87 | self.input_dim = input_dim 88 | self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False) 89 | 90 | self.use_DiffAug=True 91 | self.avg_loss=False 92 | 93 | self.cnns = nn.ModuleList() 94 | for _ in range(self.num_scales): 95 | self.cnns.append(self._make_net()) 96 | 97 | def _make_net(self): 98 | dim = self.dim 99 | cnn_x = [] 100 | cnn_x += [Conv2dBlock(self.input_dim, dim, 4, 2, 1, norm='none', activation=self.activ, pad_type=self.pad_type)] 101 | for i in range(self.n_layer - 1): 102 | cnn_x += [Conv2dBlock(dim, dim * 2, 4, 2, 1, norm=self.norm, activation=self.activ, pad_type=self.pad_type)] 103 | dim *= 2 104 | cnn_x += [nn.Conv2d(dim, 1, 4, 1, 1)] 105 | cnn_x = nn.Sequential(*cnn_x) 106 | return cnn_x 107 | 108 | def forward(self, x): 109 | if self.use_DiffAug: 110 | policy="color,translation,cutout" 111 | x = DiffAugment(x, policy=policy) 112 | outputs = [] 113 | for model in self.cnns: 114 | output=model(x) 115 | outputs.append(output) 116 | x = self.downsample(x) 117 | return outputs 118 | 119 | def calc_dis_loss(self, input_fake, input_real): 120 | # calculate the loss to train D 121 | input_real.requires_grad_() 122 | outs0 = self.forward(input_fake) 123 | outs1 = self.forward(input_real) 124 | loss = 0 125 | for it, (out0, out1) in enumerate(zip(outs0, outs1)): 126 | if self.gan_type == 'lsgan': 127 | loss += torch.mean((out0 - 0) ** 2) + torch.mean((out1 - 1) ** 2) 128 | elif self.gan_type == 'nsgan': 129 | all0 = Variable(torch.zeros_like(out0.data).cuda(), requires_grad=False) 130 | all1 = Variable(torch.ones_like(out1.data).cuda(), requires_grad=False) 131 | loss += torch.mean(F.binary_cross_entropy(F.sigmoid(out0), all0) + 132 | F.binary_cross_entropy(F.sigmoid(out1), all1)) 133 | elif self.gan_type == 'hinge': 134 | loss += F.relu(1.0+out0).mean()+F.relu(1.0-out1).mean() 135 | elif self.gan_type == 'ralsgan': 136 | loss += (torch.mean((out0 - torch. mean(out1,0) - 1) ** 2) + torch.mean((out1 - torch.mean(out0,0) + 1) ** 2))/2 137 | else: 138 | assert 0, "Unsupported GAN type: {}".format(self.gan_type) 139 | #loss+=self.r1_reg(out1,input_real) 140 | if self.avg_loss: 141 | return loss/self.num_scales 142 | return loss 143 | 144 | def calc_gen_loss(self, input_fake,input_real): 145 | # calculate the loss to train G 146 | outs0 = self.forward(input_fake) 147 | outs1 = self.forward(input_real) 148 | loss = 0 149 | for it, (out0, out1) in enumerate(zip(outs0, outs1)): 150 | if self.gan_type == 'lsgan': 151 | loss += torch.mean((out0 - 1)**2)# LSGAN 152 | elif self.gan_type == 'nsgan': 153 | all1 = Variable(torch.ones_like(out0.data).cuda(), requires_grad=False) 154 | loss += torch.mean(F.binary_cross_entropy(F.sigmoid(out0), all1)) 155 | elif self.gan_type == 'hinge': 156 | loss += -out0.mean() 157 | elif self.gan_type == 'ralsgan': 158 | loss += (torch.mean((out0 - torch.mean(out1,0) + 1) ** 2) + torch.mean((out1 - torch.mean(out0,0) - 1) ** 2))/2 159 | else: 160 | assert 0, "Unsupported GAN type: {}".format(self.gan_type) 161 | if self.avg_loss: 162 | return loss/self.num_scales 163 | return loss -------------------------------------------------------------------------------- /models/VIT.py: -------------------------------------------------------------------------------- 1 | import torch.nn.functional as F 2 | from .ops import * 3 | from collections import OrderedDict 4 | 5 | class VisionTransformer(nn.Module): 6 | """ Vision Transformer with support for patch or hybrid CNN input stage 7 | """ 8 | def __init__(self, 9 | img_size=224, 10 | patch_size=16, 11 | in_chans=3, 12 | num_classes=1000, 13 | embed_dim=768, 14 | depth=12, 15 | num_heads=12, 16 | mlp_ratio=4., 17 | qkv_bias=False, 18 | qk_scale=None, 19 | drop_rate=0., 20 | attn_drop_rate=0., 21 | drop_path_rate=0., 22 | norm_layer=nn.LayerNorm, 23 | init_values=0., 24 | use_learnable_pos_emb=False, 25 | use_rel_pos_bias=False, 26 | use_shared_rel_pos_bias=False, 27 | init_scale=0., 28 | use_mean_pooling=True): 29 | super().__init__() 30 | self.num_classes = num_classes 31 | self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models 32 | 33 | self.patch_embed = PatchEmbed( 34 | img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) 35 | num_patches = self.patch_embed.num_patches 36 | 37 | # self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) 38 | if use_learnable_pos_emb: 39 | self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) 40 | else: 41 | # sine-cosine positional embeddings is on the way 42 | self.pos_embed = get_sinusoid_encoding_table(num_patches, embed_dim) 43 | 44 | if use_shared_rel_pos_bias: 45 | self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads) 46 | else: 47 | self.rel_pos_bias = None 48 | 49 | self.pos_drop = nn.Dropout(p=drop_rate) 50 | self.use_rel_pos_bias = use_rel_pos_bias 51 | 52 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule 53 | self.blocks = nn.ModuleList([ 54 | Block( 55 | dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, 56 | drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, 57 | init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None) 58 | for i in range(depth)]) 59 | self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim) 60 | self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None 61 | self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity() 62 | 63 | if use_learnable_pos_emb: 64 | trunc_normal_(self.pos_embed, std=.02) 65 | 66 | # trunc_normal_(self.cls_token, std=.02) 67 | trunc_normal_(self.head.weight, std=.02) 68 | self.apply(self._init_weights) 69 | 70 | self.head.weight.data.mul_(init_scale) 71 | self.head.bias.data.mul_(init_scale) 72 | 73 | def _init_weights(self, m): 74 | if isinstance(m, nn.Linear): 75 | trunc_normal_(m.weight, std=.02) 76 | if isinstance(m, nn.Linear) and m.bias is not None: 77 | nn.init.constant_(m.bias, 0) 78 | elif isinstance(m, nn.LayerNorm): 79 | nn.init.constant_(m.bias, 0) 80 | nn.init.constant_(m.weight, 1.0) 81 | 82 | def get_num_layers(self): 83 | return len(self.blocks) 84 | 85 | @torch.jit.ignore 86 | def no_weight_decay(self): 87 | return {'pos_embed', 'cls_token'} 88 | 89 | def get_classifier(self): 90 | return self.head 91 | 92 | def reset_classifier(self, num_classes, global_pool=''): 93 | self.num_classes = num_classes 94 | self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() 95 | 96 | def forward_features(self, x, mask): 97 | x = self.patch_embed(x) 98 | 99 | # cls_tokens = self.cls_token.expand(batch_size, -1, -1) 100 | # x = torch.cat((cls_tokens, x), dim=1) 101 | x = x + self.pos_embed.type_as(x).to(x.device).clone().detach() 102 | 103 | B, _, C = x.shape 104 | x_vis = x[~mask].reshape(B, -1, C) # ~mask means visible 105 | 106 | rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None 107 | for blk in self.blocks: 108 | x_vis = blk(x_vis, rel_pos_bias=rel_pos_bias) 109 | 110 | # for blk in self.blocks: 111 | # x_vis = blk(x_vis) 112 | 113 | x_vis = self.norm(x_vis) 114 | return x_vis 115 | 116 | def forward_features123(self, x, return_src=False): 117 | x = self.patch_embed(x) 118 | B, _, _ = x.size() 119 | 120 | # cls_tokens = self.cls_token.expand(batch_size, -1, -1) # stole cls_tokens impl from Phil Wang, thanks 121 | # x = torch.cat((cls_tokens, x), dim=1) 122 | if self.pos_embed is not None: 123 | x = x + self.pos_embed.expand(B, -1, -1).type_as(x).to(x.device).clone().detach() 124 | x = self.pos_drop(x) 125 | 126 | for blk in self.blocks: 127 | x = blk(x) 128 | x = self.norm(x) 129 | 130 | if return_src: 131 | return x 132 | if self.fc_norm is not None: 133 | # return self.fc_norm(x[:, 1:].mean(1)) 134 | return self.fc_norm(x.mean(1)) 135 | else: 136 | return x[:, 0] 137 | 138 | def forward(self, x): 139 | x = self.forward_features(x) 140 | x = self.head(x) 141 | return x 142 | 143 | def load_state_dict(model, state_dict, prefix='', ignore_missing="relative_position_index"): 144 | missing_keys = [] 145 | unexpected_keys = [] 146 | error_msgs = [] 147 | # copy state_dict so _load_from_state_dict can modify it 148 | metadata = getattr(state_dict, '_metadata', None) 149 | state_dict = state_dict.copy() 150 | if metadata is not None: 151 | state_dict._metadata = metadata 152 | 153 | def load(module, prefix=''): 154 | local_metadata = {} if metadata is None else metadata.get( 155 | prefix[:-1], {}) 156 | module._load_from_state_dict( 157 | state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs) 158 | for name, child in module._modules.items(): 159 | if child is not None: 160 | load(child, prefix + name + '.') 161 | 162 | load(model, prefix=prefix) 163 | 164 | warn_missing_keys = [] 165 | ignore_missing_keys = [] 166 | for key in missing_keys: 167 | keep_flag = True 168 | for ignore_key in ignore_missing.split('|'): 169 | if ignore_key in key: 170 | keep_flag = False 171 | break 172 | if keep_flag: 173 | warn_missing_keys.append(key) 174 | else: 175 | ignore_missing_keys.append(key) 176 | 177 | missing_keys = warn_missing_keys 178 | 179 | if len(missing_keys) > 0: 180 | print("Weights of {} not initialized from pretrained model: {}".format( 181 | model.__class__.__name__, missing_keys)) 182 | if len(unexpected_keys) > 0: 183 | print("Weights from pretrained model not used in {}: {}".format( 184 | model.__class__.__name__, unexpected_keys)) 185 | if len(ignore_missing_keys) > 0: 186 | print("Ignored weights of {} not initialized from pretrained model: {}".format( 187 | model.__class__.__name__, ignore_missing_keys)) 188 | if len(error_msgs) > 0: 189 | print('\n'.join(error_msgs)) 190 | 191 | from functools import partial 192 | 193 | def vit_base_patch16(pretrained=False,init_ckpt=None, **kwargs): 194 | model = VisionTransformer( 195 | patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, 196 | norm_layer=partial(nn.LayerNorm, eps=1e-6), use_mean_pooling=False, **kwargs) 197 | 198 | if pretrained: 199 | checkpoint = torch.load(init_ckpt, map_location='cpu') 200 | checkpoint_model = None 201 | for model_key in 'model|module'.split('|'): 202 | if model_key in checkpoint: 203 | checkpoint_model = checkpoint[model_key] 204 | print("Load state_dict by model_key = %s" % model_key) 205 | break 206 | if checkpoint_model is None: 207 | checkpoint_model = checkpoint 208 | state_dict = model.state_dict() 209 | for k in ['head.weight', 'head.bias']: 210 | if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape: 211 | print(f"Removing key {k} from pretrained checkpoint") 212 | del checkpoint_model[k] 213 | all_keys = list(checkpoint_model.keys()) 214 | new_dict = OrderedDict() 215 | for key in all_keys: 216 | if key.startswith('backbone.'): 217 | new_dict[key[9:]] = checkpoint_model[key] 218 | elif key.startswith('encoder.'): 219 | new_dict[key[8:]] = checkpoint_model[key] 220 | else: 221 | new_dict[key] = checkpoint_model[key] 222 | checkpoint_model = new_dict 223 | if 'pos_embed' in checkpoint_model: 224 | pos_embed_checkpoint = checkpoint_model['pos_embed'] 225 | embedding_size = pos_embed_checkpoint.shape[-1] 226 | num_patches = model.patch_embed.num_patches 227 | num_extra_tokens = model.pos_embed.shape[-2] - num_patches 228 | # height (== width) for the checkpoint position embedding 229 | orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5) 230 | # height (== width) for the new position embedding 231 | new_size = int(num_patches ** 0.5) 232 | # class_token and dist_token are kept unchanged 233 | if orig_size != new_size: 234 | print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size)) 235 | extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens] 236 | # only the position tokens are interpolated 237 | pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:] 238 | pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2) 239 | pos_tokens = torch.nn.functional.interpolate( 240 | pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False) 241 | pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2) 242 | new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1) 243 | checkpoint_model['pos_embed'] = new_pos_embed 244 | 245 | load_state_dict(model, checkpoint_model, prefix='') 246 | return model 247 | 248 | if __name__=='__main__': 249 | model2=vit_base_patch16(pretrained=True,img_size=224,init_values=0.1,use_shared_rel_pos_bias=True,init_ckpt='beit_base_patch16_224_pt22k.pth') 250 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /models/ops.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from timm.models.layers import trunc_normal_ as __call_trunc_normal_ 5 | from timm.models.layers import drop_path, to_2tuple, trunc_normal_ 6 | from timm.models.registry import register_model 7 | 8 | class DropPath(nn.Module): 9 | """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). 10 | """ 11 | 12 | def __init__(self, drop_prob=None): 13 | super(DropPath, self).__init__() 14 | self.drop_prob = drop_prob 15 | 16 | def forward(self, x): 17 | return drop_path(x, self.drop_prob, self.training) 18 | 19 | def extra_repr(self) -> str: 20 | return 'p={}'.format(self.drop_prob) 21 | 22 | class Mlp(nn.Module): 23 | def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): 24 | super().__init__() 25 | out_features = out_features or in_features 26 | hidden_features = hidden_features or in_features 27 | self.fc1 = nn.Linear(in_features, hidden_features) 28 | self.act = act_layer() 29 | self.fc2 = nn.Linear(hidden_features, out_features) 30 | self.drop = nn.Dropout(drop) 31 | 32 | def forward(self, x): 33 | x = self.fc1(x) 34 | x = self.act(x) 35 | # x = self.drop(x) 36 | # commit this for the orignal BERT implement 37 | x = self.fc2(x) 38 | x = self.drop(x) 39 | return x 40 | 41 | class Attention(nn.Module): 42 | def __init__( 43 | self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., 44 | proj_drop=0., window_size=None, attn_head_dim=None): 45 | super().__init__() 46 | self.num_heads = num_heads 47 | head_dim = dim // num_heads 48 | if attn_head_dim is not None: 49 | head_dim = attn_head_dim 50 | all_head_dim = head_dim * self.num_heads 51 | self.scale = qk_scale or head_dim ** -0.5 52 | 53 | self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False) 54 | if qkv_bias: 55 | self.q_bias = nn.Parameter(torch.zeros(all_head_dim)) 56 | self.v_bias = nn.Parameter(torch.zeros(all_head_dim)) 57 | else: 58 | self.q_bias = None 59 | self.v_bias = None 60 | 61 | if window_size: 62 | self.window_size = window_size 63 | self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3 64 | self.relative_position_bias_table = nn.Parameter( 65 | torch.zeros(self.num_relative_distance, num_heads)) # 2*Wh-1 * 2*Ww-1, nH 66 | # cls to token & token 2 cls & cls to cls 67 | 68 | # get pair-wise relative position index for each token inside the window 69 | coords_h = torch.arange(window_size[0]) 70 | coords_w = torch.arange(window_size[1]) 71 | coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww 72 | coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww 73 | relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww 74 | relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 75 | relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0 76 | relative_coords[:, :, 1] += window_size[1] - 1 77 | relative_coords[:, :, 0] *= 2 * window_size[1] - 1 78 | relative_position_index = \ 79 | torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype) 80 | relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww 81 | relative_position_index[0, 0:] = self.num_relative_distance - 3 82 | relative_position_index[0:, 0] = self.num_relative_distance - 2 83 | relative_position_index[0, 0] = self.num_relative_distance - 1 84 | 85 | self.register_buffer("relative_position_index", relative_position_index) 86 | else: 87 | self.window_size = None 88 | self.relative_position_bias_table = None 89 | self.relative_position_index = None 90 | 91 | self.attn_drop = nn.Dropout(attn_drop) 92 | self.proj = nn.Linear(all_head_dim, dim) 93 | self.proj_drop = nn.Dropout(proj_drop) 94 | 95 | def forward(self, x, rel_pos_bias=None): 96 | B, N, C = x.shape 97 | qkv_bias = None 98 | if self.q_bias is not None: 99 | qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias)) 100 | # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) 101 | qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias) 102 | qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) 103 | q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) 104 | 105 | q = q * self.scale 106 | attn = (q @ k.transpose(-2, -1)) 107 | 108 | if self.relative_position_bias_table is not None: 109 | relative_position_bias = \ 110 | self.relative_position_bias_table[self.relative_position_index.view(-1)].view( 111 | self.window_size[0] * self.window_size[1] + 1, 112 | self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH 113 | relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww 114 | attn = attn + relative_position_bias.unsqueeze(0) 115 | 116 | if rel_pos_bias is not None: 117 | attn = attn + rel_pos_bias 118 | 119 | attn = attn.softmax(dim=-1) 120 | attn = self.attn_drop(attn) 121 | 122 | x = (attn @ v).transpose(1, 2).reshape(B, N, -1) 123 | x = self.proj(x) 124 | x = self.proj_drop(x) 125 | return x 126 | 127 | class CrossAttention(nn.Module): 128 | def __init__( 129 | self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., 130 | proj_drop=0., window_size=None, attn_head_dim=None): 131 | super().__init__() 132 | self.num_heads = num_heads 133 | head_dim = dim // num_heads 134 | if attn_head_dim is not None: 135 | head_dim = attn_head_dim 136 | all_head_dim = head_dim * self.num_heads 137 | self.scale = qk_scale or head_dim ** -0.5 138 | 139 | self.q = nn.Linear(dim, all_head_dim * 1, bias=False) 140 | self.kv = nn.Linear(dim, all_head_dim * 2, bias=False) 141 | 142 | if qkv_bias: 143 | self.q_bias = nn.Parameter(torch.zeros(all_head_dim)) 144 | self.v_bias = nn.Parameter(torch.zeros(all_head_dim)) 145 | else: 146 | self.q_bias = None 147 | self.v_bias = None 148 | 149 | if window_size: 150 | self.window_size = window_size 151 | self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3 152 | self.relative_position_bias_table = nn.Parameter( 153 | torch.zeros(self.num_relative_distance, num_heads)) # 2*Wh-1 * 2*Ww-1, nH 154 | # cls to token & token 2 cls & cls to cls 155 | 156 | # get pair-wise relative position index for each token inside the window 157 | coords_h = torch.arange(window_size[0]) 158 | coords_w = torch.arange(window_size[1]) 159 | coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww 160 | coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww 161 | relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww 162 | relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 163 | relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0 164 | relative_coords[:, :, 1] += window_size[1] - 1 165 | relative_coords[:, :, 0] *= 2 * window_size[1] - 1 166 | relative_position_index = \ 167 | torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype) 168 | relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww 169 | relative_position_index[0, 0:] = self.num_relative_distance - 3 170 | relative_position_index[0:, 0] = self.num_relative_distance - 2 171 | relative_position_index[0, 0] = self.num_relative_distance - 1 172 | 173 | self.register_buffer("relative_position_index", relative_position_index) 174 | else: 175 | self.window_size = None 176 | self.relative_position_bias_table = None 177 | self.relative_position_index = None 178 | 179 | self.attn_drop = nn.Dropout(attn_drop) 180 | self.proj = nn.Linear(all_head_dim, dim) 181 | self.proj_drop = nn.Dropout(proj_drop) 182 | 183 | def forward(self, x, y, rel_pos_bias=None): 184 | B, N1, C = x.shape 185 | B, N2, C = y.shape 186 | q_bias = None 187 | kv_bias = None 188 | if self.q_bias is not None: 189 | q_bias=self.q_bias 190 | kv_bias = torch.cat((torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias)) 191 | # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) 192 | q=F.linear(input=x,weight=self.q.weight,bias=q_bias) 193 | kv = F.linear(input=y, weight=self.kv.weight, bias=kv_bias) 194 | q = q.reshape(B, N1, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4) 195 | q = q[0] 196 | kv = kv.reshape(B, N2, 2, self.num_heads, -1).permute(2, 0, 3, 1, 4) 197 | k, v = kv[0], kv[1] # make torchscript happy (cannot use tensor as tuple) 198 | 199 | q = q * self.scale 200 | attn = (q @ k.transpose(-2, -1)) 201 | 202 | if self.relative_position_bias_table is not None: 203 | relative_position_bias = \ 204 | self.relative_position_bias_table[self.relative_position_index.view(-1)].view( 205 | self.window_size[0] * self.window_size[1] + 1, 206 | self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH 207 | relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww 208 | attn = attn + relative_position_bias.unsqueeze(0) 209 | 210 | if rel_pos_bias is not None: 211 | attn = attn + rel_pos_bias 212 | 213 | attn = attn.softmax(dim=-1) 214 | attn = self.attn_drop(attn) 215 | 216 | x = (attn @ v).transpose(1, 2).reshape(B, N1, -1) 217 | x = self.proj(x) 218 | x = self.proj_drop(x) 219 | return x 220 | 221 | class CorssAttnBlock(nn.Module): 222 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., 223 | drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm, 224 | window_size=None, attn_head_dim=None): 225 | super().__init__() 226 | self.norm1 = norm_layer(dim) 227 | self.attn = Attention( 228 | dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 229 | attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim) 230 | self.cross_attn = CrossAttention( 231 | dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 232 | attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim) 233 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here 234 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() 235 | self.norm2 = norm_layer(dim) 236 | self.norm3= norm_layer(dim) 237 | mlp_hidden_dim = int(dim * mlp_ratio) 238 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) 239 | 240 | if init_values > 0: 241 | self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) 242 | self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) 243 | self.gamma_3 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) 244 | else: 245 | self.gamma_1, self.gamma_2, self.gamma_3 = None, None, None 246 | 247 | def forward(self, x, y, rel_pos_bias=None): 248 | if self.gamma_1 is None: 249 | x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias)) 250 | x = x + self.drop_path(self.cross_attn(self.norm2(x),y, rel_pos_bias=rel_pos_bias)) 251 | x = x + self.drop_path(self.mlp(self.norm3(x))) 252 | else: 253 | x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias)) 254 | x = x + self.drop_path(self.gamma_2 * self.cross_attn(self.norm2(x),y, rel_pos_bias=rel_pos_bias)) 255 | x = x + self.drop_path(self.gamma_3 * self.mlp(self.norm3(x))) 256 | return x 257 | 258 | class Block(nn.Module): 259 | 260 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., 261 | drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm, 262 | window_size=None, attn_head_dim=None): 263 | super().__init__() 264 | self.norm1 = norm_layer(dim) 265 | self.attn = Attention( 266 | dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 267 | attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim) 268 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here 269 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() 270 | self.norm2 = norm_layer(dim) 271 | mlp_hidden_dim = int(dim * mlp_ratio) 272 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) 273 | 274 | if init_values > 0: 275 | self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) 276 | self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) 277 | else: 278 | self.gamma_1, self.gamma_2 = None, None 279 | 280 | def forward(self, x, rel_pos_bias=None): 281 | if self.gamma_1 is None: 282 | x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias)) 283 | x = x + self.drop_path(self.mlp(self.norm2(x))) 284 | else: 285 | x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias)) 286 | x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x))) 287 | return x 288 | 289 | class PatchEmbed(nn.Module): 290 | """ Image to Patch Embedding 291 | """ 292 | 293 | def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): 294 | super().__init__() 295 | img_size = to_2tuple(img_size) 296 | patch_size = to_2tuple(patch_size) 297 | num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) 298 | self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) 299 | self.img_size = img_size 300 | self.patch_size = patch_size 301 | self.num_patches = num_patches 302 | 303 | self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) 304 | 305 | def forward(self, x, **kwargs): 306 | B, C, H, W = x.shape 307 | # FIXME look at relaxing size constraints 308 | assert H == self.img_size[0] and W == self.img_size[1], \ 309 | f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." 310 | x = self.proj(x).flatten(2).transpose(1, 2) 311 | return x 312 | 313 | class RelativePositionBias(nn.Module): 314 | def __init__(self, window_size, num_heads): 315 | super().__init__() 316 | self.window_size = window_size 317 | self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3 318 | self.relative_position_bias_table = nn.Parameter( 319 | torch.zeros(self.num_relative_distance, num_heads)) # 2*Wh-1 * 2*Ww-1, nH 320 | # cls to token & token 2 cls & cls to cls 321 | 322 | # get pair-wise relative position index for each token inside the window 323 | coords_h = torch.arange(window_size[0]) 324 | coords_w = torch.arange(window_size[1]) 325 | coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww 326 | coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww 327 | relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww 328 | relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 329 | relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0 330 | relative_coords[:, :, 1] += window_size[1] - 1 331 | relative_coords[:, :, 0] *= 2 * window_size[1] - 1 332 | relative_position_index = \ 333 | torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype) 334 | relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww 335 | relative_position_index[0, 0:] = self.num_relative_distance - 3 336 | relative_position_index[0:, 0] = self.num_relative_distance - 2 337 | relative_position_index[0, 0] = self.num_relative_distance - 1 338 | 339 | self.register_buffer("relative_position_index", relative_position_index) 340 | 341 | # trunc_normal_(self.relative_position_bias_table, std=.02) 342 | 343 | def forward(self): 344 | relative_position_bias = \ 345 | self.relative_position_bias_table[self.relative_position_index.view(-1)].view( 346 | self.window_size[0] * self.window_size[1] + 1, 347 | self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH 348 | return relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww 349 | 350 | import numpy as np 351 | def get_sinusoid_encoding_table(n_position, d_hid): 352 | ''' Sinusoid position encoding table ''' 353 | # TODO: make it with torch instead of numpy 354 | def get_position_angle_vec(position): 355 | return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] 356 | 357 | sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) 358 | sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i 359 | sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 360 | 361 | return torch.FloatTensor(sinusoid_table).unsqueeze(0) --------------------------------------------------------------------------------