dsipts 1.1.5__py3-none-any.whl

dsipts/models/d3vae/neural_operations.py

@@ -0,0 +1,314 @@
+# -*-Encoding: utf-8 -*-
+"""
+Authors:
+    Li,Yan (liyan22021121@gmail.com)
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.batchnorm import _BatchNorm
+from collections import OrderedDict
+
+
+BN_EPS = 1e-5
+SYNC_BN = False
+
+OPS = OrderedDict([
+    ('res_elu', lambda Cin, Cout, stride: ELUConv(Cin, Cout, 3, stride, 1)),
+    ('res_bnelu', lambda Cin, Cout, stride: BNELUConv(Cin, Cout, 3, stride, 1)),
+    ('res_bnswish', lambda Cin, Cout, stride: BNSwishConv(Cin, Cout, 3, stride, 1)),
+    ('res_bnswish5', lambda Cin, Cout, stride: BNSwishConv(Cin, Cout, 3, stride, 2, 2)),
+    ('mconv_e6k5g0', lambda Cin, Cout, stride: InvertedResidual(Cin, Cout, stride, ex=6, dil=1, k=5, g=1)),
+    ('mconv_e3k5g0', lambda Cin, Cout, stride: InvertedResidual(Cin, Cout, stride, ex=3, dil=1, k=5, g=1)),
+    ('mconv_e3k5g8', lambda Cin, Cout, stride: InvertedResidual(Cin, Cout, stride, ex=3, dil=1, k=5, g=8)),
+    ('mconv_e6k11g0', lambda Cin, Cout, stride: InvertedResidual(Cin, Cout, stride, ex=6, dil=1, k=11, g=0)),
+])
+
+
+class SyncBatchNormSwish(_BatchNorm):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
+                 track_running_stats=True, process_group=None):
+        super(SyncBatchNormSwish, self).__init__(num_features, eps, momentum, affine, track_running_stats)
+        self.process_group = process_group
+        self.ddp_gpu_size = None
+
+    def forward(self, input):
+        exponential_average_factor = self.momentum
+        out = F.batch_norm(
+            input, self.running_mean, self.running_var, self.weight, self.bias,
+            self.training or not self.track_running_stats,
+            exponential_average_factor, self.eps)
+        return out
+
+
+def get_skip_connection(C, stride, channel_mult):
+    if stride == 1:
+        return Identity()
+    elif stride == 2:
+        return FactorizedReduce(C, int(channel_mult * C))
+    elif stride == -1:
+        return nn.Sequential(UpSample(), Conv2D(C, int(C / channel_mult), kernel_size=1))
+
+
+def norm(t, dim):
+    return torch.sqrt(torch.sum(t * t, dim))
+
+
+def logit(t):
+    return torch.log(t) - torch.log(1 - t)
+
+
+def act(t):
+    # The following implementation has lower memory.
+    return SwishFN.apply(t)
+
+
+class SwishFN(torch.autograd.Function):
+    def forward(ctx, i):
+        result = i * torch.sigmoid(i)
+        ctx.save_for_backward(i)
+        return result
+
+    def backward(ctx, grad_output):
+        i = ctx.saved_variables[0]
+        sigmoid_i = torch.sigmoid(i)
+        return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i)))
+
+
+class Swish(nn.Module):
+    def __init__(self):
+        super(Swish, self).__init__()
+
+    def forward(self, x):
+        return act(x)
+
+
+def normalize_weight_jit(log_weight_norm, weight):
+    n = torch.exp(log_weight_norm)
+    wn = torch.sqrt(torch.sum(weight * weight, dim=[1, 2, 3]))   # norm(w)
+    weight = n * weight / (wn.view(-1, 1, 1, 1) + 1e-5)
+    return weight
+
+
+class Conv2D(nn.Conv2d):
+    """Allows for weights as input."""
+
+    def __init__(self, C_in, C_out, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=False, data_init=False,
+                 weight_norm=True):
+        """
+        Args:
+            use_shared (bool): Use weights for this layer or not?
+        """
+        super(Conv2D, self).__init__(C_in, C_out, kernel_size, stride, padding, dilation, groups, bias)
+
+        self.log_weight_norm = None
+        if weight_norm:
+            init = norm(self.weight, dim=[1, 2, 3]).view(-1, 1, 1, 1)
+            self.log_weight_norm = nn.Parameter(torch.log(init + 1e-2), requires_grad=True)
+
+        self.data_init = data_init
+        self.init_done = False
+        self.weight_normalized = self.normalize_weight()
+
+    def forward(self, x):
+        # do data based initialization
+        self.weight_normalized = self.normalize_weight()
+        #print(self.weight_normalized.shape)
+        bias = self.bias
+        return F.conv2d(x, self.weight_normalized, bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+
+    def normalize_weight(self):
+        """ applies weight normalization """
+        if self.log_weight_norm is not None:
+            weight = normalize_weight_jit(self.log_weight_norm, self.weight)
+        else:
+            weight = self.weight
+
+        return weight
+
+
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+
+
+class SyncBatchNorm(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(SyncBatchNorm, self).__init__()
+        self.bn = nn.BatchNorm(*args, **kwargs)
+
+    def forward(self, x):
+        return self.bn(x)
+
+
+# quick switch between multi-gpu, single-gpu batch norm
+def get_batchnorm(*args, **kwargs):
+    return nn.BatchNorm2d(*args, **kwargs)
+
+
+class ELUConv(nn.Module):
+    def __init__(self, C_in, C_out, kernel_size, stride=1, padding=0, dilation=1):
+        super(ELUConv, self).__init__()
+        self.upsample = stride == -1
+        stride = abs(stride)
+        self.conv_0 = Conv2D(C_in, C_out, kernel_size, stride=stride, padding=padding, bias=True, dilation=dilation,
+                             data_init=True)
+
+    def forward(self, x):
+        out = F.elu(x)
+        if self.upsample:
+            out = F.interpolate(out, scale_factor=2, mode='nearest')
+        out = self.conv_0(out)
+        return out
+
+
+class BNELUConv(nn.Module):
+    def __init__(self, C_in, C_out, kernel_size, stride=1, padding=0, dilation=1):
+        super(BNELUConv, self).__init__()
+        self.upsample = stride == -1
+        stride = abs(stride)
+        self.bn = get_batchnorm(C_in, eps=BN_EPS, momentum=0.05)
+        self.conv_0 = Conv2D(C_in, C_out, kernel_size, stride=stride, padding=padding, bias=True, dilation=dilation)
+
+    def forward(self, x):
+        x = self.bn(x)
+        out = F.elu(x)
+        if self.upsample:
+            out = F.interpolate(out, scale_factor=2, mode='nearest')
+        out = self.conv_0(out)
+        return out
+
+
+class BNSwishConv(nn.Module):
+    """ReLU + Conv2d + BN."""
+
+    def __init__(self, C_in, C_out, kernel_size, stride=1, padding=0, dilation=1):
+        super(BNSwishConv, self).__init__()
+        self.upsample = stride == -1
+        stride = abs(stride)
+        self.bn_act = SyncBatchNormSwish(C_in, eps=BN_EPS, momentum=0.05)
+        self.conv_0 = Conv2D(C_in, C_out, kernel_size, stride=stride, padding=padding, bias=True, dilation=dilation)
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.Tensor): of size (B, C_in, H, W)
+        """
+        out = self.bn_act(x)
+        if self.upsample:
+            out = F.interpolate(out, scale_factor=2, mode='nearest')
+        out = self.conv_0(out)
+        return out
+
+
+class FactorizedReduce(nn.Module):
+    def __init__(self, C_in, C_out):
+        super(FactorizedReduce, self).__init__()
+        assert C_out % 2 == 0
+        self.conv_1 = Conv2D(C_in, C_out // 4, 1, stride=2, padding=0, bias=True)
+        self.conv_2 = Conv2D(C_in, C_out // 4, 1, stride=2, padding=0, bias=True)
+        self.conv_3 = Conv2D(C_in, C_out // 4, 1, stride=2, padding=0, bias=True)
+        self.conv_4 = Conv2D(C_in, C_out - 3 * (C_out // 4), 1, stride=2, padding=0, bias=True)
+
+    def forward(self, x):
+        out = act(x)
+        conv1 = self.conv_1(out[:,:,:, :])
+        conv2 = self.conv_2(out[:, :, 1:, :])
+        conv3 = self.conv_3(out[:, :, :, :])
+        conv4 = self.conv_4(out[:, :, 1:, :])
+        out = torch.cat([conv1, conv2, conv3, conv4], dim=1)
+        return out
+
+
+class UpSample(nn.Module):
+    def __init__(self):
+        super(UpSample, self).__init__()
+        pass
+
+    def forward(self, x):
+        return F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
+
+
+class EncCombinerCell(nn.Module):
+    def __init__(self, Cin1, Cin2, Cout, cell_type):
+        super(EncCombinerCell, self).__init__()
+        self.cell_type = cell_type
+        # Cin = Cin1 + Cin2
+        self.conv = Conv2D(Cin2, Cout, kernel_size=1, stride=1, padding=0, bias=True)
+
+    def forward(self, x1, x2):
+        x2 = self.conv(x2)
+        out = x1 + x2
+        return out
+
+
+# original combiner
+class DecCombinerCell(nn.Module):
+    def __init__(self, Cin1, Cin2, Cout, cell_type):
+        super(DecCombinerCell, self).__init__()
+        self.cell_type = cell_type
+        self.conv = Conv2D(Cin1 + Cin2, Cout, kernel_size=1, stride=1, padding=0, bias=True)
+
+    def forward(self, x1, x2):
+        out = torch.cat([x1, x2], dim=1)
+        out = self.conv(out)
+        return out
+
+
+class ConvBNSwish(nn.Module):
+    def __init__(self, Cin, Cout, k=3, stride=1, groups=1, dilation=1):
+        padding = dilation * (k - 1) // 2
+        super(ConvBNSwish, self).__init__()
+
+        self.conv = nn.Sequential(
+            Conv2D(Cin, Cout, k, stride, padding, groups=groups, bias=False, dilation=dilation, weight_norm=False),
+            SyncBatchNormSwish(Cout, eps=BN_EPS, momentum=0.05)  # drop in replacement for BN + Swish
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class SE(nn.Module):
+    def __init__(self, Cin, Cout):
+        super(SE, self).__init__()
+        num_hidden = max(Cout // 16, 4)
+        self.se = nn.Sequential(nn.Linear(Cin, num_hidden), nn.ReLU(inplace=True),
+                                nn.Linear(num_hidden, Cout), nn.Sigmoid())
+
+    def forward(self, x):
+        se = torch.mean(x, dim=[2, 3])
+        se = se.view(se.size(0), -1)
+        se = self.se(se)
+        se = se.view(se.size(0), -1, 1, 1)
+        return x * se
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, Cin, Cout, stride, ex, dil, k, g):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2, -1]
+
+        hidden_dim = int(round(Cin * ex))
+        self.use_res_connect = self.stride == 1 and Cin == Cout
+        self.upsample = self.stride == -1
+        self.stride = abs(self.stride)
+        groups = hidden_dim if g == 0 else g
+
+        layers0 = [nn.UpsamplingNearest2d(scale_factor=2)] if self.upsample else []
+        layers = [get_batchnorm(Cin, eps=BN_EPS, momentum=0.05),
+                  ConvBNSwish(Cin, hidden_dim, k=1),
+                  ConvBNSwish(hidden_dim, hidden_dim, stride=self.stride, groups=groups, k=k, dilation=dil),
+                  Conv2D(hidden_dim, Cout, 1, 1, 0, bias=False, weight_norm=False),
+                  get_batchnorm(Cout, momentum=0.05)]
+
+        layers0.extend(layers)
+        self.conv = nn.Sequential(*layers0)
+
+    def forward(self, x):
+        return self.conv(x)

dsipts/models/d3vae/resnet.py

@@ -0,0 +1,153 @@
+# -*-Encoding: utf-8 -*-
+"""
+Authors:
+    Li,Yan (liyan22021121@gmail.com)
+"""
+from torch import nn
+
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(0.0, 0.2)
+    elif classname.find("BatchNorm") != -1:
+        m.weight.data.normal_(1.0, 0.2)
+        m.bias.data.fill_(0)
+
+
+class MyConvo2d(nn.Module):
+    def __init__(self, input_dim, output_dim, kernel_size,  stride = 1, bias = True):
+        super(MyConvo2d, self).__init__()
+        self.padding = int((kernel_size - 1)/2)
+        self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride=1, padding=self.padding, bias = bias)
+
+    def forward(self, input):
+        output = self.conv(input)
+        return output
+
+
+class Square(nn.Module):
+    def __init__(self):
+        super(Square,self).__init__()
+        pass
+    
+    def forward(self,in_vect):
+        return in_vect**2
+
+
+class Swish(nn.Module):
+    def __init__(self):
+        super(Swish,self).__init__()
+        pass
+    
+    def forward(self,in_vect):
+        return in_vect*nn.functional.sigmoid(in_vect)
+
+
+class MeanPoolConv(nn.Module):
+    def __init__(self, input_dim, output_dim, kernel_size):
+        super(MeanPoolConv, self).__init__()
+        self.conv = MyConvo2d(input_dim, output_dim, kernel_size)
+
+    def forward(self, input):
+        output = input
+        output = self.conv(output)
+        return output
+
+
+class ConvMeanPool(nn.Module):
+    def __init__(self, input_dim, output_dim, kernel_size):
+        super(ConvMeanPool, self).__init__()
+        self.conv = MyConvo2d(input_dim, output_dim, kernel_size)
+
+    def forward(self, input):
+        output = self.conv(input)
+        return output
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, input_dim, output_dim, kernel_size, hw, resample=None, normalize=False,AF=nn.ELU()):
+        super(ResidualBlock, self).__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.kernel_size = kernel_size
+        self.resample = resample
+        self.normalize = normalize
+        self.bn1 = None
+        self.bn2 = None
+        self.relu1 = AF
+        self.relu2 = AF
+        if resample == 'down':
+            self.bn1 = nn.LayerNorm([input_dim, hw, hw])
+            self.bn2 = nn.LayerNorm([input_dim, hw, hw])
+        elif resample == 'none':
+            self.bn1 = nn.LayerNorm([input_dim, hw, hw])
+            self.bn2 = nn.LayerNorm([input_dim, hw, hw])
+
+        if resample == 'down':
+            self.conv_shortcut = MeanPoolConv(input_dim, output_dim, kernel_size = 1)
+            self.conv_1 = MyConvo2d(input_dim, input_dim, kernel_size = kernel_size, bias = False)
+            self.conv_2 = ConvMeanPool(input_dim, output_dim, kernel_size = kernel_size)
+        elif resample == 'none':
+            self.conv_shortcut = MyConvo2d(input_dim, output_dim, kernel_size = 1)
+            self.conv_1 = MyConvo2d(input_dim, input_dim, kernel_size = kernel_size, bias = False)
+            self.conv_2 = MyConvo2d(input_dim, output_dim, kernel_size = kernel_size)
+                    
+    def forward(self, input):
+        if self.input_dim == self.output_dim and self.resample is None:
+                shortcut = input
+        else:
+            shortcut = self.conv_shortcut(input)
+        
+        if self.normalize is False:
+            output = input
+            output = self.relu1(output)
+            output = self.conv_1(output)
+            output = self.relu2(output)
+            output = self.conv_2(output)
+        else:
+            output = input
+            output = self.bn1(output)
+            output = self.relu1(output)
+            output = self.conv_1(output)
+            output = self.bn2(output)
+            output = self.relu2(output)
+            output = self.conv_2(output)
+
+        return shortcut + output
+
+
+class Res12_Quadratic(nn.Module):
+    def __init__(self,inchan,dim,hw,normalize=False,AF=None):
+        super(Res12_Quadratic, self).__init__()
+        
+        self.hw = hw
+        self.dim = dim
+        self.inchan = inchan
+        self.conv1 = MyConvo2d(inchan,dim, 3)
+        self.rb1 = ResidualBlock(dim, 2*dim, 3, int(hw), resample = 'down',normalize=normalize,AF=AF)
+        self.rbc1 = ResidualBlock(2*dim, 2*dim, 3, int(hw/2), resample = 'none',normalize=normalize,AF=AF)
+        self.rb2 = ResidualBlock(2*dim, 4*dim, 3, int(hw/2), resample = 'down',normalize=normalize,AF=AF)
+        self.rbc2 = ResidualBlock(4*dim, 4*dim, 3, int(hw/4), resample = 'none',normalize=normalize,AF=AF)
+        self.rb3 = ResidualBlock(4*dim, 8*dim, 3, int(hw/4), resample = 'down',normalize=normalize,AF=AF)
+        self.rbc3 = ResidualBlock(8*dim, 8*dim, 3, int(hw/8), resample = 'none',normalize=normalize,AF=AF)
+        self.ln1 = nn.Linear(int(hw/8)*int(hw/8)*8*dim, 1)
+        self.ln2 = nn.Linear(int(hw/8)*int(hw/8)*8*dim, 1)
+        self.lq = nn.Linear(int(hw/8)*int(hw/8)*8*dim, 1)
+        self.Square = Square()
+                
+    def forward(self, x_in):
+        output = x_in
+        output = self.conv1(output)
+        # print(output.shape)
+        output = self.rb1(output)
+        output = self.rbc1(output)
+        output = self.rb2(output)
+        output = self.rbc2(output)
+        output = self.rb3(output)
+        output = self.rbc3(output)
+        output = output.view(-1, int(self.hw/8)*int(self.hw/8)*8*self.dim)
+        output = self.ln1(output)*self.ln2(output)+self.lq(self.Square(output))
+        output = output.view(-1)
+        return output