Spaces:
Build error
Build error
| import imp | |
| import torch | |
| import torch.nn as nn | |
| from timm.models.layers import trunc_normal_, DropPath, to_2tuple | |
| import os | |
| from model.blocks import Mlp | |
| class query_Attention(nn.Module): | |
| def __init__(self, dim, num_heads=2, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| head_dim = dim // num_heads | |
| # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights | |
| self.scale = qk_scale or head_dim ** -0.5 | |
| self.q = nn.Parameter(torch.ones((1, 10, dim)), requires_grad=True) | |
| self.k = nn.Linear(dim, dim, bias=qkv_bias) | |
| self.v = nn.Linear(dim, dim, 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 | |
| k = self.k(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) | |
| v = self.v(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) | |
| q = self.q.expand(B, -1, -1).view(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) | |
| # k = self.k(x).reshape(B, N, self.num_heads, torch.div(C,self.num_heads, rounding_mode='floor')).permute(0, 2, 1, 3) | |
| # v = self.v(x).reshape(B, N, self.num_heads, torch.div(C,self.num_heads, rounding_mode='floor')).permute(0, 2, 1, 3) | |
| # q = self.q.expand(B, -1, -1).view(B, -1, self.num_heads, torch.div(C,self.num_heads, rounding_mode='floor')).permute(0, 2, 1, 3) | |
| 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, 10, C) | |
| x = self.proj(x) | |
| x = self.proj_drop(x) | |
| return x | |
| class query_SABlock(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.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim) | |
| self.norm1 = norm_layer(dim) | |
| self.attn = query_Attention( | |
| dim, | |
| num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, | |
| attn_drop=attn_drop, proj_drop=drop) | |
| # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here | |
| 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): | |
| x = x + self.pos_embed(x) | |
| x = x.flatten(2).transpose(1, 2) | |
| x = self.drop_path(self.attn(self.norm1(x))) | |
| x = x + self.drop_path(self.mlp(self.norm2(x))) | |
| return x | |
| class conv_embedding(nn.Module): | |
| def __init__(self, in_channels, out_channels): | |
| super(conv_embedding, self).__init__() | |
| self.proj = nn.Sequential( | |
| nn.Conv2d(in_channels, out_channels // 2, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)), | |
| nn.BatchNorm2d(out_channels // 2), | |
| nn.GELU(), | |
| # nn.Conv2d(out_channels // 2, out_channels // 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)), | |
| # nn.BatchNorm2d(out_channels // 2), | |
| # nn.GELU(), | |
| nn.Conv2d(out_channels // 2, out_channels, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)), | |
| nn.BatchNorm2d(out_channels), | |
| ) | |
| def forward(self, x): | |
| x = self.proj(x) | |
| return x | |
| class Global_pred(nn.Module): | |
| def __init__(self, in_channels=3, out_channels=64, num_heads=4, type='exp'): | |
| super(Global_pred, self).__init__() | |
| if type == 'exp': | |
| self.gamma_base = nn.Parameter(torch.ones((1)), requires_grad=False) # False in exposure correction | |
| else: | |
| self.gamma_base = nn.Parameter(torch.ones((1)), requires_grad=True) | |
| self.color_base = nn.Parameter(torch.eye((3)), requires_grad=True) # basic color matrix | |
| # main blocks | |
| self.conv_large = conv_embedding(in_channels, out_channels) | |
| self.generator = query_SABlock(dim=out_channels, num_heads=num_heads) | |
| self.gamma_linear = nn.Linear(out_channels, 1) | |
| self.color_linear = nn.Linear(out_channels, 1) | |
| self.apply(self._init_weights) | |
| for name, p in self.named_parameters(): | |
| if name == 'generator.attn.v.weight': | |
| nn.init.constant_(p, 0) | |
| 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 forward(self, x): | |
| #print(self.gamma_base) | |
| x = self.conv_large(x) | |
| x = self.generator(x) | |
| gamma, color = x[:, 0].unsqueeze(1), x[:, 1:] | |
| gamma = self.gamma_linear(gamma).squeeze(-1) + self.gamma_base | |
| #print(self.gamma_base, self.gamma_linear(gamma)) | |
| color = self.color_linear(color).squeeze(-1).view(-1, 3, 3) + self.color_base | |
| return gamma, color | |
| if __name__ == "__main__": | |
| os.environ['CUDA_VISIBLE_DEVICES']='3' | |
| #net = Local_pred_new().cuda() | |
| img = torch.Tensor(8, 3, 400, 600) | |
| global_net = Global_pred() | |
| gamma, color = global_net(img) | |
| print(gamma.shape, color.shape) | |