graft-pytorch 0.1.7__py3-none-any.whl

graft/models/efficientnet.py

@@ -0,0 +1,197 @@
+'''EfficientNet in PyTorch.
+
+Reference
+    EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks
+    https://github.com/keras-team/keras-applications/blob/master/keras_applications/efficientnet.py
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def swish(x):
+    return x * x.sigmoid()
+
+
+def drop_connect(x, drop_ratio):
+    keep_ratio = 1.0 - drop_ratio
+    mask = torch.empty([x.shape[0], 1, 1, 1], dtype=x.dtype, device=x.device)
+    mask.bernoulli_(keep_ratio)
+    x.div_(keep_ratio)
+    x.mul_(mask)
+    return x
+
+
+class SE(nn.Module):
+    '''Squeeze-and-Excitation block with Swish.'''
+
+    def __init__(self, in_channels, se_channels):
+        super(SE, self).__init__()
+        self.se1 = nn.Conv2d(in_channels, se_channels,
+                             kernel_size=1, bias=True)
+        self.se2 = nn.Conv2d(se_channels, in_channels,
+                             kernel_size=1, bias=True)
+
+
+    def forward(self, x):
+        out = F.adaptive_avg_pool2d(x, (1, 1))
+        out = swish(self.se1(out))
+        out = self.se2(out).sigmoid()
+        out = x * out
+        return out
+
+
+class Block(nn.Module):
+    '''expansion + depthwise + pointwise + squeeze-excitation'''
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 expand_ratio=1,
+                 se_ratio=0.,
+                 drop_rate=0.):
+        super(Block, self).__init__()
+        self.stride = stride
+        self.drop_rate = drop_rate
+        self.expand_ratio = expand_ratio
+
+        # Expansion
+        channels = expand_ratio * in_channels
+        self.conv1 = nn.Conv2d(in_channels,
+                               channels,
+                               kernel_size=1,
+                               stride=1,
+                               padding=0,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(channels)
+
+        # Depthwise conv
+        self.conv2 = nn.Conv2d(channels,
+                               channels,
+                               kernel_size=kernel_size,
+                               stride=stride,
+                               padding=(1 if kernel_size == 3 else 2),
+                               groups=channels,
+                               bias=False)
+        self.bn2 = nn.BatchNorm2d(channels)
+
+        # SE layers
+        se_channels = int(in_channels * se_ratio)
+        self.se = SE(channels, se_channels)
+
+        # Output
+        self.conv3 = nn.Conv2d(channels,
+                               out_channels,
+                               kernel_size=1,
+                               stride=1,
+                               padding=0,
+                               bias=False)
+        self.bn3 = nn.BatchNorm2d(out_channels)
+
+        # Skip connection if in and out shapes are the same (MV-V2 style)
+        self.has_skip = (stride == 1) and (in_channels == out_channels)
+
+
+    def forward(self, x):
+        out = x if self.expand_ratio == 1 else swish(self.bn1(self.conv1(x)))
+        out = swish(self.bn2(self.conv2(out)))
+        out = self.se(out)
+        out = self.bn3(self.conv3(out))
+        if self.has_skip:
+            if self.training and self.drop_rate > 0:
+                out = drop_connect(out, self.drop_rate)
+            out = out + x
+        return out
+
+
+class EfficientNet(nn.Module):
+    def __init__(self, cfg, num_classes=10):
+        super(EfficientNet, self).__init__()
+        self.cfg = cfg
+        self.conv1 = nn.Conv2d(3,
+                               32,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1,
+                               bias=False)
+        self.embDim = cfg['out_channels'][-1]
+        self.bn1 = nn.BatchNorm2d(32)
+        self.layers = self._make_layers(in_channels=32)
+        self.linear = nn.Linear(cfg['out_channels'][-1], num_classes)
+
+
+    def _make_layers(self, in_channels):
+        layers = []
+        cfg = [self.cfg[k] for k in ['expansion', 'out_channels', 'num_blocks', 'kernel_size',
+                                     'stride']]
+        b = 0
+        blocks = sum(self.cfg['num_blocks'])
+        for expansion, out_channels, num_blocks, kernel_size, stride in zip(*cfg):
+            strides = [stride] + [1] * (num_blocks - 1)
+            for stride in strides:
+                drop_rate = self.cfg['drop_connect_rate'] * b / blocks
+                layers.append(
+                    Block(in_channels,
+                          out_channels,
+                          kernel_size,
+                          stride,
+                          expansion,
+                          se_ratio=0.25,
+                          drop_rate=drop_rate))
+                in_channels = out_channels
+        return nn.Sequential(*layers)
+
+
+    def forward(self, x, last=False, freeze=False):
+        if freeze:
+            with torch.no_grad():
+                out = swish(self.bn1(self.conv1(x)))
+                out = self.layers(out)
+                out = F.adaptive_avg_pool2d(out, 1)
+                e = out.view(out.size(0), -1)
+                dropout_rate = self.cfg['dropout_rate']
+                if self.training and dropout_rate > 0:
+                    e = F.dropout(e, p=dropout_rate)
+        else:
+            out = swish(self.bn1(self.conv1(x)))
+            out = self.layers(out)
+            out = F.adaptive_avg_pool2d(out, 1)
+            e = out.view(out.size(0), -1)
+            dropout_rate = self.cfg['dropout_rate']
+            if self.training and dropout_rate > 0:
+                e = F.dropout(e, p=dropout_rate)
+        out = self.linear(e)
+        if last:
+            return out, e
+        else:
+            return out
+
+
+    def get_embedding_dim(self):
+        return self.embDim
+
+
+def EfficientNetB0(num_classes=10):
+    cfg = {
+        'num_blocks': [1, 2, 2, 3, 3, 4, 1],
+        'expansion': [1, 6, 6, 6, 6, 6, 6],
+        'out_channels': [16, 24, 40, 80, 112, 192, 320],
+        'kernel_size': [3, 3, 5, 3, 5, 5, 3],
+        'stride': [1, 2, 2, 2, 1, 2, 1],
+        'dropout_rate': 0.2,
+        'drop_connect_rate': 0.2,
+    }
+    return EfficientNet(cfg, num_classes)
+
+
+def test():
+    net = EfficientNetB0()
+    x = torch.randn(2, 3, 32, 32)
+    y = net(x)
+    print(y.shape)
+
+
+if __name__ == '__main__':
+    test()

graft/models/efficientnetb7.py

@@ -0,0 +1,268 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+import math
+
+class Swish(nn.Module):
+    def __init__(self):
+        super(Swish, self).__init__()
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        return x * self.sigmoid(x)
+
+def _RoundChannels(c, divisor=8, min_value=None):
+    if min_value is None:
+        min_value = divisor
+    new_c = max(min_value, int(c + divisor / 2) // divisor * divisor)
+    if new_c < 0.9 * c:
+        new_c += divisor
+    return new_c
+
+def _RoundRepeats(r):
+    return int(math.ceil(r))
+
+def _DropPath(x, drop_prob, training):
+    if drop_prob > 0 and training:
+        keep_prob = 1 - drop_prob
+        if x.is_cuda:
+            mask = Variable(torch.cuda.FloatTensor(x.size(0), 1, 1, 1).bernoulli_(keep_prob))
+        else:
+            mask = Variable(torch.FloatTensor(x.size(0), 1, 1, 1).bernoulli_(keep_prob))
+        x.div_(keep_prob)
+        x.mul_(mask)
+
+    return x
+
+def _BatchNorm(channels, eps=1e-3, momentum=0.01):
+    return nn.BatchNorm2d(channels, eps=eps, momentum=momentum)
+
+def _Conv3x3Bn(in_channels, out_channels, stride):
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False),
+        _BatchNorm(out_channels),
+        Swish()
+    )
+
+def _Conv1x1Bn(in_channels, out_channels):
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 1, 1, 0, bias=False),
+        _BatchNorm(out_channels),
+        Swish()
+    )
+
+class SqueezeAndExcite(nn.Module):
+    def __init__(self, channels, squeeze_channels, se_ratio):
+        super(SqueezeAndExcite, self).__init__()
+
+        squeeze_channels = squeeze_channels * se_ratio
+        if not squeeze_channels.is_integer():
+            raise ValueError('channels must be divisible by 1/ratio')
+
+        squeeze_channels = int(squeeze_channels)
+        self.se_reduce = nn.Conv2d(channels, squeeze_channels, 1, 1, 0, bias=True)
+        self.non_linear1 = Swish()
+        self.se_expand = nn.Conv2d(squeeze_channels, channels, 1, 1, 0, bias=True)
+        self.non_linear2 = nn.Sigmoid()
+
+    def forward(self, x):
+        y = torch.mean(x, (2, 3), keepdim=True)
+        y = self.non_linear1(self.se_reduce(y))
+        y = self.non_linear2(self.se_expand(y))
+        y = x * y
+
+        return y
+
+class MBConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio, se_ratio, drop_path_rate):
+        super(MBConvBlock, self).__init__()
+
+        expand = (expand_ratio != 1)
+        expand_channels = in_channels * expand_ratio
+        se = (se_ratio != 0.0)
+        self.residual_connection = (stride == 1 and in_channels == out_channels)
+        self.drop_path_rate = drop_path_rate
+
+        conv = []
+
+        if expand:
+            # expansion phase
+            pw_expansion = nn.Sequential(
+                nn.Conv2d(in_channels, expand_channels, 1, 1, 0, bias=False),
+                _BatchNorm(expand_channels),
+                Swish()
+            )
+            conv.append(pw_expansion)
+
+        # depthwise convolution phase
+        dw = nn.Sequential(
+            nn.Conv2d(
+                expand_channels,
+                expand_channels,
+                kernel_size,
+                stride,
+                kernel_size//2,
+                groups=expand_channels,
+                bias=False
+            ),
+            _BatchNorm(expand_channels),
+            Swish()
+        )
+        conv.append(dw)
+
+        if se:
+            # squeeze and excite
+            squeeze_excite = SqueezeAndExcite(expand_channels, in_channels, se_ratio)
+            conv.append(squeeze_excite)
+
+        # projection phase
+        pw_projection = nn.Sequential(
+            nn.Conv2d(expand_channels, out_channels, 1, 1, 0, bias=False),
+            _BatchNorm(out_channels)
+        )
+        conv.append(pw_projection)
+
+        self.conv = nn.Sequential(*conv)
+
+    def forward(self, x):
+        if self.residual_connection:
+            return x + _DropPath(self.conv(x), self.drop_path_rate, self.training)
+        else:
+            return self.conv(x)
+
+class EfficientNet(nn.Module):
+    config = [
+        #(in_channels, out_channels, kernel_size, stride, expand_ratio, se_ratio, repeats)
+        [32,  16,  3, 1, 1, 0.25, 1],
+        [16,  24,  3, 2, 6, 0.25, 2],
+        [24,  40,  5, 2, 6, 0.25, 2],
+        [40,  80,  3, 2, 6, 0.25, 3],
+        [80,  112, 5, 1, 6, 0.25, 3],
+        [112, 192, 5, 2, 6, 0.25, 4],
+        [192, 320, 3, 1, 6, 0.25, 1]
+    ]
+
+    def __init__(self, param, num_classes=100, stem_channels=32, feature_size=1280, drop_connect_rate=0.2):
+        super(EfficientNet, self).__init__()
+
+        # scaling width
+        width_coefficient = param[0]
+        if width_coefficient != 1.0:
+            stem_channels = _RoundChannels(stem_channels*width_coefficient)
+            for conf in self.config:
+                conf[0] = _RoundChannels(conf[0]*width_coefficient)
+                conf[1] = _RoundChannels(conf[1]*width_coefficient)
+
+        # scaling depth
+        depth_coefficient = param[1]
+        if depth_coefficient != 1.0:
+            for conf in self.config:
+                conf[6] = _RoundRepeats(conf[6]*depth_coefficient)
+
+        # scaling resolution
+        input_size = param[2]
+
+        # stem convolution
+        self.stem_conv = _Conv3x3Bn(3, stem_channels, 2)
+
+        # total #blocks
+        total_blocks = 0
+        for conf in self.config:
+            total_blocks += conf[6]
+
+        # mobile inverted bottleneck
+        blocks = []
+        for in_channels, out_channels, kernel_size, stride, expand_ratio, se_ratio, repeats in self.config:
+            # drop connect rate based on block index
+            drop_rate = drop_connect_rate * (len(blocks) / total_blocks)
+            blocks.append(MBConvBlock(in_channels, out_channels, kernel_size, stride, expand_ratio, se_ratio, drop_rate))
+            for _ in range(repeats-1):
+                drop_rate = drop_connect_rate * (len(blocks) / total_blocks)
+                blocks.append(MBConvBlock(out_channels, out_channels, kernel_size, 1, expand_ratio, se_ratio, drop_rate))
+        self.blocks = nn.Sequential(*blocks)
+
+        # last several layers
+        self.head_conv = _Conv1x1Bn(self.config[-1][1], feature_size)
+        #self.avgpool = nn.AvgPool2d(input_size//32, stride=1)
+        self.dropout = nn.Dropout(param[3])
+        self.classifier = nn.Linear(feature_size, num_classes)
+
+        self._initialize_weights()
+
+    def forward(self, x, last=False, freeze=False):
+            if freeze:
+                with torch.no_grad():
+                    x = self.stem_conv(x)
+                    x = self.blocks(x)
+                    x = self.head_conv(x)
+                    #x = self.avgpool(x)
+                    #x = x.view(x.size(0), -1)
+                    x = torch.mean(x, (2, 3))
+                    e = self.dropout(x)
+                    # x = self.classifier(x)
+            else:
+                x = self.stem_conv(x)
+                x = self.blocks(x)
+                x = self.head_conv(x)
+                #x = self.avgpool(x)
+                #x = x.view(x.size(0), -1)
+                x = torch.mean(x, (2, 3))
+                e = self.dropout(x)
+            out = self.classifier(x)
+            if last:
+                return out, e
+            else:
+                return out
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2.0 / n))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                n = m.weight.size(1)
+                m.weight.data.normal_(0, 0.01)
+                m.bias.data.zero_()
+
+def EfficientNetB7(num_classes=100):
+    net_param = {
+        # 'efficientnet type': (width_coef, depth_coef, resolution, dropout_rate)
+        'efficientnet-b0': (1.0, 1.0, 224, 0.2),
+        'efficientnet-b1': (1.0, 1.1, 240, 0.2),
+        'efficientnet-b2': (1.1, 1.2, 260, 0.3),
+        'efficientnet-b3': (1.2, 1.4, 300, 0.3),
+        'efficientnet-b4': (1.4, 1.8, 380, 0.4),
+        'efficientnet-b5': (1.6, 2.2, 456, 0.4),
+        'efficientnet-b6': (1.8, 2.6, 528, 0.5),
+        'efficientnet-b7': (2.0, 3.1, 600, 0.5)
+    }
+
+    param = net_param['efficientnet-b7']
+    #net = EfficientNet(param)
+    return EfficientNet(param,  num_classes)
+
+if __name__ == '__main__':
+    net_param = {
+        # 'efficientnet type': (width_coef, depth_coef, resolution, dropout_rate)
+        'efficientnet-b0': (1.0, 1.0, 224, 0.2),
+        'efficientnet-b1': (1.0, 1.1, 240, 0.2),
+        'efficientnet-b2': (1.1, 1.2, 260, 0.3),
+        'efficientnet-b3': (1.2, 1.4, 300, 0.3),
+        'efficientnet-b4': (1.4, 1.8, 380, 0.4),
+        'efficientnet-b5': (1.6, 2.2, 456, 0.4),
+        'efficientnet-b6': (1.8, 2.6, 528, 0.5),
+        'efficientnet-b7': (2.0, 3.1, 600, 0.5)
+    }
+
+    param = net_param['efficientnet-b0']
+    net = EfficientNet(param)
+    x_image = Variable(torch.randn(1, 3, param[2], param[2]))
+    y = net(x_image)

graft/models/fashioncnn.py

@@ -0,0 +1,69 @@
+import torch.nn as nn
+import torch
+from torch import Tensor
+
+
+class FashionCNN(nn.Module):
+    
+    def __init__(self,in_channels, num_classes):
+        super(FashionCNN, self).__init__()
+        
+        self.layer1 = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(2)
+        )
+        
+        self.fc1 = nn.Linear(in_features=64*6*6, out_features=600)
+        self.drop = nn.Dropout2d(0.25)
+        self.fc2 = nn.Linear(in_features=600, out_features=120)
+        self.fc3 = nn.Linear(in_features=120, out_features=num_classes)
+        
+#     def forward(self, x):
+#         out = self.layer1(x)
+#         out = self.layer2(out)
+#         out = out.view(out.size(0), -1)
+#         out = self.fc1(out)
+#         out = self.drop(out)
+#         out = self.fc2(out)
+#         out = self.fc3(out)
+        
+#         return out
+    
+    def forward(self, x: Tensor, last=False, freeze=False) -> Tensor:
+        # See note [TorchScript super()]
+        if freeze:
+            with torch.no_grad():
+                self.eval()
+                x = self.layer1(x)
+                x = self.layer2(x)
+                x = x.view(x.size(0), -1)
+                x = self.fc1(x)
+                x = self.drop(x)
+                x = self.fc2(x)
+#                 x = self.fc3(x)
+                features = torch.flatten(x, 1)
+                self.train()
+        else:
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = x.view(x.size(0), -1)
+            x = self.fc1(x)
+            x = self.drop(x)
+            x = self.fc2(x)
+#             x = self.fc3(x)
+            features = torch.flatten(x, 1)
+
+        out = self.fc3(features)
+        if last:
+            return out, features
+        else:
+            return out

graft/models/mobilenet.py

@@ -0,0 +1,83 @@
+'''MobileNet in PyTorch.
+
+Reference
+    MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications 
+    https://arxiv.org/abs/1704.04861
+'''
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Block(nn.Module):
+    '''Depthwise conv + Pointwise conv'''
+    
+    def __init__(self, in_planes, out_planes, stride=1):
+        super(Block, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, in_planes, kernel_size=3, stride=stride, padding=1, groups=in_planes, bias=False)
+        self.bn1 = nn.BatchNorm2d(in_planes)
+        self.conv2 = nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_planes)
+
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        return out
+
+
+class MobileNet(nn.Module):
+    # (128,2) means conv planes=128, conv stride=2, by default conv stride=1
+    cfg = [64, (128,2), 128, (256,2), 256, (512,2), 512, 512, 512, 512, 512, (1024,2), 1024]
+
+    def __init__(self, num_classes=10):
+        super(MobileNet, self).__init__()
+        self.embDim = 1024
+        
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.layers = self._make_layers(in_planes=32)
+        self.linear = nn.Linear(1024, num_classes)
+
+
+    def _make_layers(self, in_planes):
+        layers = []
+        for x in self.cfg:
+            out_planes = x if isinstance(x, int) else x[0]
+            stride = 1 if isinstance(x, int) else x[1]
+            layers.append(Block(in_planes, out_planes, stride))
+            in_planes = out_planes
+        return nn.Sequential(*layers)
+
+
+    def forward(self, x, last=False, freeze=False):
+        if freeze:
+            with torch.no_grad():
+                out = F.relu(self.bn1(self.conv1(x)))
+                out = self.layers(out)
+                out = F.avg_pool2d(out, 2)
+                e = out.view(out.size(0), -1)
+        else:
+            out = F.relu(self.bn1(self.conv1(x)))
+            out = self.layers(out)
+            out = F.avg_pool2d(out, 2)
+            e = out.view(out.size(0), -1)
+        out = self.linear(e)
+        if last:
+            return out, e
+        else:
+            return out
+
+    def get_embedding_dim(self):
+        return self.embDim
+
+        
+def test():
+    net = MobileNet()
+    x = torch.randn(1, 3, 32, 32)
+    y = net(x)
+    print(y.size())
+
+# test()