Spaces:
Running
on
L40S
Running
on
L40S
| import torch.nn as nn | |
| import torch | |
| import math | |
| import torch.nn.functional as F | |
| class single_conv(nn.Module): | |
| def __init__(self, in_ch, out_ch): | |
| super(single_conv, self).__init__() | |
| self.conv = nn.Sequential( | |
| nn.Conv2d(in_ch, out_ch, 3, stride=1, padding=1), | |
| nn.BatchNorm2d(out_ch), | |
| nn.ReLU(inplace=True), | |
| ) | |
| def forward(self, x): | |
| return self.conv(x) | |
| class double_conv(nn.Module): | |
| def __init__(self, in_ch, out_ch): | |
| super(double_conv, self).__init__() | |
| self.conv = nn.Sequential( | |
| nn.Conv2d(in_ch, out_ch, 3, stride=1, padding=1), nn.BatchNorm2d(out_ch), | |
| nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, stride=1, padding=1), | |
| nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) | |
| ) | |
| def forward(self, x): | |
| return self.conv(x) | |
| class double_conv_down(nn.Module): | |
| def __init__(self, in_ch, out_ch): | |
| super(double_conv_down, self).__init__() | |
| self.conv = nn.Sequential( | |
| nn.Conv2d(in_ch, out_ch, 3, stride=2, padding=1), nn.BatchNorm2d(out_ch), | |
| nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, stride=1, padding=1), | |
| nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) | |
| ) | |
| def forward(self, x): | |
| return self.conv(x) | |
| class double_conv_up(nn.Module): | |
| def __init__(self, in_ch, out_ch): | |
| super(double_conv_up, self).__init__() | |
| self.conv = nn.Sequential( | |
| nn.UpsamplingNearest2d(scale_factor=2), | |
| nn.Conv2d(in_ch, out_ch, 3, stride=1, padding=1), nn.BatchNorm2d(out_ch), | |
| nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, stride=1, padding=1), | |
| nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) | |
| ) | |
| def forward(self, x): | |
| return self.conv(x) | |
| class PosEnSine(nn.Module): | |
| """ | |
| Code borrowed from DETR: models/positional_encoding.py | |
| output size: b*(2.num_pos_feats)*h*w | |
| """ | |
| def __init__(self, num_pos_feats): | |
| super(PosEnSine, self).__init__() | |
| self.num_pos_feats = num_pos_feats | |
| self.normalize = True | |
| self.scale = 2 * math.pi | |
| self.temperature = 10000 | |
| def forward(self, x, pt_coord=None): | |
| b, c, h, w = x.shape | |
| if pt_coord is not None: | |
| z_embed = pt_coord[:, :, 2].unsqueeze(-1) + 1. | |
| y_embed = pt_coord[:, :, 1].unsqueeze(-1) + 1. | |
| x_embed = pt_coord[:, :, 0].unsqueeze(-1) + 1. | |
| else: | |
| not_mask = torch.ones(1, h, w, device=x.device) | |
| y_embed = not_mask.cumsum(1, dtype=torch.float32) | |
| x_embed = not_mask.cumsum(2, dtype=torch.float32) | |
| z_embed = torch.ones_like(x_embed) | |
| if self.normalize: | |
| eps = 1e-6 | |
| z_embed = z_embed / (torch.max(z_embed) + eps) * self.scale | |
| y_embed = y_embed / (torch.max(y_embed) + eps) * self.scale | |
| x_embed = x_embed / (torch.max(x_embed) + eps) * self.scale | |
| dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) | |
| dim_t = self.temperature**(2 * (dim_t // 2) / self.num_pos_feats) | |
| pos_x = x_embed[:, :, :, None] / dim_t | |
| pos_y = y_embed[:, :, :, None] / dim_t | |
| pos_z = z_embed[:, :, :, None] / dim_t | |
| pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), | |
| dim=4).flatten(3) | |
| pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), | |
| dim=4).flatten(3) | |
| pos_z = torch.stack((pos_z[:, :, :, 0::2].sin(), pos_z[:, :, :, 1::2].cos()), | |
| dim=4).flatten(3) | |
| pos = torch.cat((pos_x, pos_y, pos_z), dim=3).permute(0, 3, 1, 2) | |
| # if pt_coord is None: | |
| pos = pos.repeat(b, 1, 1, 1) | |
| return pos | |
| def softmax_attention(q, k, v): | |
| # b x n x d x h x w | |
| h, w = q.shape[-2], q.shape[-1] | |
| q = q.flatten(-2).transpose(-2, -1) # b x n x hw x d | |
| k = k.flatten(-2) # b x n x d x hw | |
| v = v.flatten(-2).transpose(-2, -1) | |
| print('softmax', q.shape, k.shape, v.shape) | |
| N = k.shape[-1] # ?????? maybe change to k.shape[-2]???? | |
| attn = torch.matmul(q / N**0.5, k) | |
| attn = F.softmax(attn, dim=-1) | |
| output = torch.matmul(attn, v) | |
| output = output.transpose(-2, -1) | |
| output = output.view(*output.shape[:-1], h, w) | |
| return output, attn | |
| def dotproduct_attention(q, k, v): | |
| # b x n x d x h x w | |
| h, w = q.shape[-2], q.shape[-1] | |
| q = q.flatten(-2).transpose(-2, -1) # b x n x hw x d | |
| k = k.flatten(-2) # b x n x d x hw | |
| v = v.flatten(-2).transpose(-2, -1) | |
| N = k.shape[-1] | |
| attn = None | |
| tmp = torch.matmul(k, v) / N | |
| output = torch.matmul(q, tmp) | |
| output = output.transpose(-2, -1) | |
| output = output.view(*output.shape[:-1], h, w) | |
| return output, attn | |
| def long_range_attention(q, k, v, P_h, P_w): # fixed patch size | |
| B, N, C, qH, qW = q.size() | |
| _, _, _, kH, kW = k.size() | |
| qQ_h, qQ_w = qH // P_h, qW // P_w | |
| kQ_h, kQ_w = kH // P_h, kW // P_w | |
| q = q.reshape(B, N, C, qQ_h, P_h, qQ_w, P_w) | |
| k = k.reshape(B, N, C, kQ_h, P_h, kQ_w, P_w) | |
| v = v.reshape(B, N, -1, kQ_h, P_h, kQ_w, P_w) | |
| q = q.permute(0, 1, 4, 6, 2, 3, 5) # [b, n, Ph, Pw, d, Qh, Qw] | |
| k = k.permute(0, 1, 4, 6, 2, 3, 5) | |
| v = v.permute(0, 1, 4, 6, 2, 3, 5) | |
| output, attn = softmax_attention(q, k, v) # attn: [b, n, Ph, Pw, qQh*qQw, kQ_h*kQ_w] | |
| output = output.permute(0, 1, 4, 5, 2, 6, 3) | |
| output = output.reshape(B, N, -1, qH, qW) | |
| return output, attn | |
| def short_range_attention(q, k, v, Q_h, Q_w): # fixed patch number | |
| B, N, C, qH, qW = q.size() | |
| _, _, _, kH, kW = k.size() | |
| qP_h, qP_w = qH // Q_h, qW // Q_w | |
| kP_h, kP_w = kH // Q_h, kW // Q_w | |
| q = q.reshape(B, N, C, Q_h, qP_h, Q_w, qP_w) | |
| k = k.reshape(B, N, C, Q_h, kP_h, Q_w, kP_w) | |
| v = v.reshape(B, N, -1, Q_h, kP_h, Q_w, kP_w) | |
| q = q.permute(0, 1, 3, 5, 2, 4, 6) # [b, n, Qh, Qw, d, Ph, Pw] | |
| k = k.permute(0, 1, 3, 5, 2, 4, 6) | |
| v = v.permute(0, 1, 3, 5, 2, 4, 6) | |
| output, attn = softmax_attention(q, k, v) # attn: [b, n, Qh, Qw, qPh*qPw, kPh*kPw] | |
| output = output.permute(0, 1, 4, 2, 5, 3, 6) | |
| output = output.reshape(B, N, -1, qH, qW) | |
| return output, attn | |
| def space_to_depth(x, block_size): | |
| x_shape = x.shape | |
| c, h, w = x_shape[-3:] | |
| if len(x.shape) >= 5: | |
| x = x.view(-1, c, h, w) | |
| unfolded_x = torch.nn.functional.unfold(x, block_size, stride=block_size) | |
| return unfolded_x.view(*x_shape[0:-3], c * block_size**2, h // block_size, w // block_size) | |
| def depth_to_space(x, block_size): | |
| x_shape = x.shape | |
| c, h, w = x_shape[-3:] | |
| x = x.view(-1, c, h, w) | |
| y = torch.nn.functional.pixel_shuffle(x, block_size) | |
| return y.view(*x_shape[0:-3], -1, h * block_size, w * block_size) | |
| def patch_attention(q, k, v, P): | |
| # q: [b, nhead, c, h, w] | |
| q_patch = space_to_depth(q, P) # [b, nhead, cP^2, h/P, w/P] | |
| k_patch = space_to_depth(k, P) | |
| v_patch = space_to_depth(v, P) | |
| # output: [b, nhead, cP^2, h/P, w/P] | |
| # attn: [b, nhead, h/P*w/P, h/P*w/P] | |
| output, attn = softmax_attention(q_patch, k_patch, v_patch) | |
| output = depth_to_space(output, P) # output: [b, nhead, c, h, w] | |
| return output, attn | |