ViT.py

#from dino https://github.com/facebookresearch/dino
import math
import torch
import torch.nn as nn
import warnings
import os,sys
import numpy as np
import torch.nn.functional as F
from tqdm import tqdm
get_path = os.path.dirname(os.path.abspath(__file__))
sys.path.append(get_path)
current_path = os.path.dirname(os.path.abspath(__file__)).split('/')

def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)

def drop_path(x, drop_prob: float = 0., training: bool = False):
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x, attn


class Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, return_attention=False):
        y, attn = self.attn(self.norm1(x))
        if return_attention:
            return attn
        x = x + self.drop_path(y)
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=[224,224], patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        num_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x


class RBFExpansion(nn.Module):
    r"""Expand distances between nodes by radial basis functions.

    .. math::
        \exp(- \gamma * ||d - \mu||^2)

    """
    def __init__(self, low=0., high=1., num_centers=265):
        super(RBFExpansion, self).__init__()

        # num_centers = int(np.ceil((high - low) / gap))
        centers = np.linspace(low, high, num_centers).reshape(1,-1)
        self.register_buffer('centers',torch.tensor(centers).float())
        self.coef = num_centers / (low - high)
        self.linear = nn.Linear(num_centers,num_centers)

    def forward(self, values): 
        #input shape [B]
        radial = values[:,None] - self.centers
        x =  torch.exp(self.coef * (radial ** 2))
        x = self.linear(x)
        return x

class resblock(nn.Module):

    def __init__(
        self,
        in_channel=256,
        hidden=512,
        out_channel = 3
    ) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(in_channel,hidden,1)
        self.norm1 = nn.LayerNorm(hidden)
        self.conv2 = nn.Conv2d(hidden,hidden,3,padding=1,padding_mode='replicate')
        self.norm2 = nn.LayerNorm(hidden)
        self.conv3 = nn.Conv2d(hidden,out_channel,1)

    def forward(self, x):

        out = self.conv1(x)
        out = out.permute(0,2,3,1).contiguous() # N C W H -> N W H C
        out = self.norm1(out)
        out = F.gelu(out)
        out = out.permute(0,3,1,2).contiguous()  # N W H C -> N C W H
        out = self.conv2(out)
        out = out.permute(0,2,3,1).contiguous()
        out = self.norm2(out)
        out = F.gelu(out)
        out = out.permute(0,3,1,2).contiguous()
        out = self.conv3(out)

        return out

class U_Vit(nn.Module):
    """ Vision Transformer """
    def __init__(self, img_size=[64,64], patch_size=4, in_chans=3, embed_dim=384, depth=7,
                 num_heads=4, mlp_ratio=1., qkv_bias=True, qk_scale=None, drop_rate=0.1, attn_drop_rate=0.1,
                 drop_path_rate=0.1, norm_layer=nn.LayerNorm,emb=PatchEmbed,cat_num=0, **kwargs):
        super().__init__()
        assert depth % 2 == 1, "U-Vit has to have odd layer number!"
        assert embed_dim % patch_size**2 == 0, 'Embed dim should be exactly divided by patch size square'
        self.depth = depth
        self.num_features = self.embed_dim = embed_dim
        self.patch_size = patch_size
        self.in_chans = in_chans
        self.img_size = img_size
        self.patch_embed = emb(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches
        self.time_embed = RBFExpansion(num_centers=embed_dim)
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 2, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.blocks = nn.ModuleList([
            Block(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
            for i in range(depth)])
        self.skip_connection =  nn.ModuleList([
            nn.Linear(2*embed_dim,embed_dim)
            for i in range(depth//2)
        ])
        self.norm = norm_layer(embed_dim)
        self.head = resblock(in_channel=embed_dim//patch_size**2,hidden=embed_dim,out_channel=in_chans) 
        self.class_emb = nn.Embedding(cat_num+1,embed_dim)
        self.cat_num = cat_num
        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.class_emb.weight, std=.02)
        trunc_normal_(self.time_embed.linear.weight, std=.02)
        nn.init.constant_(self.time_embed.linear.bias, 0)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)


    def prepare_tokens(self, x,steps,labels):
        B, nc, w, h = x.shape
        x = self.patch_embed(x)  # patch linear embedding
        #cls_tokens = self.cls_token.expand(B, -1, -1)
        time_step = self.time_embed(steps).unsqueeze_(1)
        class_token = self.class_emb(labels).unsqueeze_(1)
        x = torch.cat((class_token,time_step, x), dim=1)
        x = x + self.pos_embed 
        return self.pos_drop(x)

    def forward(self,noisy_img,t,labels):
        x1 = self.prepare_tokens(noisy_img,t,labels)
        intermidiate = []
        for i,blk in enumerate(self.blocks):
            if  i > (self.depth-1)//2:
                x1 = self.skip_connection[self.depth-1-i](torch.cat([intermidiate[self.depth-1-i],x1],-1))
            x1 = blk(x1)
            if i < (self.depth-1)//2:
                intermidiate.append(x1)
        x1 = self.norm(x1)
        velocity = x1[:,2:,:]
        velocity = velocity.view(-1,self.img_size[0]//self.patch_size,self.img_size[1]//self.patch_size,self.patch_size,self.patch_size,self.embed_dim // self.patch_size**2)
        velocity = velocity.permute(0, 5, 1, 3, 2, 4).contiguous()
        velocity = velocity.view(-1,self.embed_dim // self.patch_size**2,self.img_size[0],self.img_size[1])
        velocity = self.head(velocity)
        return velocity

    @torch.no_grad()
    def forward_transforms(self,input_samples,device,k=25,N=128,labels=None):
            # #debug
        # import matplotlib.pyplot as plt
        # import torchvision.transforms.functional as transF
        if labels is None:
            labels = torch.zeros(N).long().to(device)
        img = input_samples
        for i in tqdm(range(k)):
            t = i/k*torch.ones(N,device=device)
            img += self.forward(img,t,labels)/k
        return img,labels.cpu().numpy()

    @torch.no_grad()
    def reverse_transforms(self,input_samples,device,k=25,N=128,labels=None):
            # #debug
        # import matplotlib.pyplot as plt
        # import torchvision.transforms.functional as transF
        if labels is None:
            labels = torch.zeros(N).long().to(device)
        img = input_samples
        for i in tqdm(range(k,0,-1)):
            t = i/k*torch.ones(N,device=device)
            img -= self.forward(img,t,labels)/k
        return img,labels.cpu().numpy()

class ModelWapper4ODE(nn.Module):
    def __init__(self,model) -> None:
        super().__init__()
        self.model = model
    def forward(self,t,y):
        N,_,_,_ = y.shape
        t = t*torch.ones(N,device=y.device).float()
        labels = torch.zeros(N,device=y.device).long()
        return self.model(y,t,labels)

def RFlow_Loss(pred,img,noise):
    loss = F.huber_loss(pred,img-noise)
    return loss

if __name__ == '__main__':
    from rflow_loader import Interpolation_Dataset
    from torch.utils.data import DataLoader
    dataset = Interpolation_Dataset("/data/protein/OxfordFlowers/train",[64,64],None)
    dataloader = DataLoader(dataset,batch_size=4,shuffle=True)
    noisy_img,img,noise,t,labels = next(iter(dataloader))
    noisy_img = noisy_img.float()
    img = img.float()
    noise = noise.float()
    t = t.float()
    labels = labels.long()
    model = U_Vit()
    out = model(noisy_img,t,labels)
    loss = RFlow_Loss(out,img,noise)
    print(noise.shape)