Delete vietocr/model

vietocr/model/backbone/cnn.py

@@ -1,28 +0,0 @@
-import torch
-from torch import nn
-
-import vietocr.model.backbone.vgg as vgg
-from vietocr.model.backbone.resnet import Resnet50
-
-class CNN(nn.Module):
-    def __init__(self, backbone, **kwargs):
-        super(CNN, self).__init__()
-
-        if backbone == 'vgg11_bn':
-            self.model = vgg.vgg11_bn(**kwargs)
-        elif backbone == 'vgg19_bn':
-            self.model = vgg.vgg19_bn(**kwargs)
-        elif backbone == 'resnet50':
-            self.model = Resnet50(**kwargs)
-
-    def forward(self, x):
-        return self.model(x)
-
-    def freeze(self):
-        for name, param in self.model.features.named_parameters():
-            if name != 'last_conv_1x1':
-                param.requires_grad = False
-
-    def unfreeze(self):
-        for param in self.model.features.parameters():
-            param.requires_grad = True

vietocr/model/backbone/resnet.py

@@ -1,140 +0,0 @@
-import torch
-from torch import nn
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(BasicBlock, self).__init__()
-        self.conv1 = self._conv3x3(inplanes, planes)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.conv2 = self._conv3x3(planes, planes)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def _conv3x3(self, in_planes, out_planes, stride=1):
-        "3x3 convolution with padding"
-        return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                         padding=1, bias=False)
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-        out += residual
-        out = self.relu(out)
-
-        return out
-    
-class ResNet(nn.Module):
-
-    def __init__(self, input_channel, output_channel, block, layers):
-        super(ResNet, self).__init__()
-
-        self.output_channel_block = [int(output_channel / 4), int(output_channel / 2), output_channel, output_channel]
-
-        self.inplanes = int(output_channel / 8)
-        self.conv0_1 = nn.Conv2d(input_channel, int(output_channel / 16),
-                                 kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn0_1 = nn.BatchNorm2d(int(output_channel / 16))
-        self.conv0_2 = nn.Conv2d(int(output_channel / 16), self.inplanes,
-                                 kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn0_2 = nn.BatchNorm2d(self.inplanes)
-        self.relu = nn.ReLU(inplace=True)
-
-        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
-        self.layer1 = self._make_layer(block, self.output_channel_block[0], layers[0])
-        self.conv1 = nn.Conv2d(self.output_channel_block[0], self.output_channel_block[
-                               0], kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(self.output_channel_block[0])
-
-        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
-        self.layer2 = self._make_layer(block, self.output_channel_block[1], layers[1], stride=1)
-        self.conv2 = nn.Conv2d(self.output_channel_block[1], self.output_channel_block[
-                               1], kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(self.output_channel_block[1])
-
-        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=(2, 1), padding=(0, 1))
-        self.layer3 = self._make_layer(block, self.output_channel_block[2], layers[2], stride=1)
-        self.conv3 = nn.Conv2d(self.output_channel_block[2], self.output_channel_block[
-                               2], kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn3 = nn.BatchNorm2d(self.output_channel_block[2])
-
-        self.layer4 = self._make_layer(block, self.output_channel_block[3], layers[3], stride=1)
-        self.conv4_1 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[
-                                 3], kernel_size=2, stride=(2, 1), padding=(0, 1), bias=False)
-        self.bn4_1 = nn.BatchNorm2d(self.output_channel_block[3])
-        self.conv4_2 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[
-                                 3], kernel_size=2, stride=1, padding=0, bias=False)
-        self.bn4_2 = nn.BatchNorm2d(self.output_channel_block[3])
-
-    def _make_layer(self, block, planes, blocks, stride=1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(planes * block.expansion),
-            )
-
-        layers = []
-        layers.append(block(self.inplanes, planes, stride, downsample))
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = self.conv0_1(x)
-        x = self.bn0_1(x)
-        x = self.relu(x)
-        x = self.conv0_2(x)
-        x = self.bn0_2(x)
-        x = self.relu(x)
-
-        x = self.maxpool1(x)
-        x = self.layer1(x)
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.relu(x)
-
-        x = self.maxpool2(x)
-        x = self.layer2(x)
-        x = self.conv2(x)
-        x = self.bn2(x)
-        x = self.relu(x)
-
-        x = self.maxpool3(x)
-        x = self.layer3(x)
-        x = self.conv3(x)
-        x = self.bn3(x)
-        x = self.relu(x)
-
-        x = self.layer4(x)
-        x = self.conv4_1(x)
-        x = self.bn4_1(x)
-        x = self.relu(x)
-        x = self.conv4_2(x)
-        x = self.bn4_2(x)
-        conv = self.relu(x)
-        
-        conv = conv.transpose(-1, -2)
-        conv = conv.flatten(2)
-        conv = conv.permute(-1, 0, 1)
-
-        return conv
-
-def Resnet50(ss, hidden):
-    return ResNet(3, hidden, BasicBlock, [1, 2, 5, 3])
-

vietocr/model/backbone/vgg.py

@@ -1,50 +0,0 @@
-import torch
-from torch import nn
-from torchvision import models
-from einops import rearrange
-from torchvision.models._utils import IntermediateLayerGetter
-
-
-class Vgg(nn.Module):
-    def __init__(self, name, ss, ks, hidden, pretrained=True, dropout=0.5):
-        super(Vgg, self).__init__()
-
-        if name == 'vgg11_bn':
-            cnn = models.vgg11_bn(weights='DEFAULT')
-        elif name == 'vgg19_bn':
-            cnn = models.vgg19_bn(weights='DEFAULT')
-
-        pool_idx = 0
-        
-        for i, layer in enumerate(cnn.features):
-            if isinstance(layer, torch.nn.MaxPool2d):        
-                cnn.features[i] = torch.nn.AvgPool2d(kernel_size=ks[pool_idx], stride=ss[pool_idx], padding=0)
-                pool_idx += 1
- 
-        self.features = cnn.features
-        self.dropout = nn.Dropout(dropout)
-        self.last_conv_1x1 = nn.Conv2d(512, hidden, 1)
-
-    def forward(self, x):
-        """
-        Shape: 
-            - x: (N, C, H, W)
-            - output: (W, N, C)
-        """
-
-        conv = self.features(x)
-        conv = self.dropout(conv)
-        conv = self.last_conv_1x1(conv)
-
-#        conv = rearrange(conv, 'b d h w -> b d (w h)')
-        conv = conv.transpose(-1, -2)
-        conv = conv.flatten(2)
-        conv = conv.permute(-1, 0, 1)
-        return conv
-
-def vgg11_bn(ss, ks, hidden, pretrained=True, dropout=0.5):
-    return Vgg('vgg11_bn', ss, ks, hidden, pretrained, dropout)
-
-def vgg19_bn(ss, ks, hidden, pretrained=True, dropout=0.5):
-    return Vgg('vgg19_bn', ss, ks, hidden, pretrained, dropout)
-   

vietocr/model/seqmodel/convseq2seq.py

@@ -1,324 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torch.nn.functional as F
-
-class Encoder(nn.Module):
-    def __init__(self, 
-                 emb_dim, 
-                 hid_dim, 
-                 n_layers, 
-                 kernel_size, 
-                 dropout, 
-                 device,
-                 max_length = 512):
-        super().__init__()
-        
-        assert kernel_size % 2 == 1, "Kernel size must be odd!"
-        
-        self.device = device
-        
-        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device)
-        
-#         self.tok_embedding = nn.Embedding(input_dim, emb_dim)
-        self.pos_embedding = nn.Embedding(max_length, emb_dim)
-        
-        self.emb2hid = nn.Linear(emb_dim, hid_dim)
-        self.hid2emb = nn.Linear(hid_dim, emb_dim)
-        
-        self.convs = nn.ModuleList([nn.Conv1d(in_channels = hid_dim, 
-                                              out_channels = 2 * hid_dim, 
-                                              kernel_size = kernel_size, 
-                                              padding = (kernel_size - 1) // 2)
-                                    for _ in range(n_layers)])
-        
-        self.dropout = nn.Dropout(dropout)
-        
-    def forward(self, src):
-        
-        #src = [batch size, src len]
-        
-        src = src.transpose(0, 1)
-        
-        batch_size = src.shape[0]
-        src_len = src.shape[1]
-        device = src.device
-        
-        #create position tensor
-        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(device)
-        
-        #pos = [0, 1, 2, 3, ..., src len - 1]
-        
-        #pos = [batch size, src len]
-        
-        #embed tokens and positions
-        
-#         tok_embedded = self.tok_embedding(src)
-        tok_embedded = src
-    
-        pos_embedded = self.pos_embedding(pos)
-        
-        #tok_embedded = pos_embedded = [batch size, src len, emb dim]
-        
-        #combine embeddings by elementwise summing
-        embedded = self.dropout(tok_embedded + pos_embedded)
-        
-        #embedded = [batch size, src len, emb dim]
-        
-        #pass embedded through linear layer to convert from emb dim to hid dim
-        conv_input = self.emb2hid(embedded)
-        
-        #conv_input = [batch size, src len, hid dim]
-        
-        #permute for convolutional layer
-        conv_input = conv_input.permute(0, 2, 1) 
-        
-        #conv_input = [batch size, hid dim, src len]
-        
-        #begin convolutional blocks...
-        
-        for i, conv in enumerate(self.convs):
-        
-            #pass through convolutional layer
-            conved = conv(self.dropout(conv_input))
-
-            #conved = [batch size, 2 * hid dim, src len]
-
-            #pass through GLU activation function
-            conved = F.glu(conved, dim = 1)
-
-            #conved = [batch size, hid dim, src len]
-            
-            #apply residual connection
-            conved = (conved + conv_input) * self.scale
-
-            #conved = [batch size, hid dim, src len]
-            
-            #set conv_input to conved for next loop iteration
-            conv_input = conved
-        
-        #...end convolutional blocks
-        
-        #permute and convert back to emb dim
-        conved = self.hid2emb(conved.permute(0, 2, 1))
-        
-        #conved = [batch size, src len, emb dim]
-        
-        #elementwise sum output (conved) and input (embedded) to be used for attention
-        combined = (conved + embedded) * self.scale
-        
-        #combined = [batch size, src len, emb dim]
-        
-        return conved, combined
-
-class Decoder(nn.Module):
-    def __init__(self, 
-                 output_dim, 
-                 emb_dim, 
-                 hid_dim, 
-                 n_layers, 
-                 kernel_size, 
-                 dropout, 
-                 trg_pad_idx, 
-                 device,
-                 max_length = 512):
-        super().__init__()
-        
-        self.kernel_size = kernel_size
-        self.trg_pad_idx = trg_pad_idx
-        self.device = device
-        
-        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device)
-        
-        self.tok_embedding = nn.Embedding(output_dim, emb_dim)
-        self.pos_embedding = nn.Embedding(max_length, emb_dim)
-        
-        self.emb2hid = nn.Linear(emb_dim, hid_dim)
-        self.hid2emb = nn.Linear(hid_dim, emb_dim)
-        
-        self.attn_hid2emb = nn.Linear(hid_dim, emb_dim)
-        self.attn_emb2hid = nn.Linear(emb_dim, hid_dim)
-        
-        self.fc_out = nn.Linear(emb_dim, output_dim)
-        
-        self.convs = nn.ModuleList([nn.Conv1d(in_channels = hid_dim, 
-                                              out_channels = 2 * hid_dim, 
-                                              kernel_size = kernel_size)
-                                    for _ in range(n_layers)])
-        
-        self.dropout = nn.Dropout(dropout)
-      
-    def calculate_attention(self, embedded, conved, encoder_conved, encoder_combined):
-        
-        #embedded = [batch size, trg len, emb dim]
-        #conved = [batch size, hid dim, trg len]
-        #encoder_conved = encoder_combined = [batch size, src len, emb dim]
-        
-        #permute and convert back to emb dim
-        conved_emb = self.attn_hid2emb(conved.permute(0, 2, 1))
-        
-        #conved_emb = [batch size, trg len, emb dim]
-        
-        combined = (conved_emb + embedded) * self.scale
-        
-        #combined = [batch size, trg len, emb dim]
-                
-        energy = torch.matmul(combined, encoder_conved.permute(0, 2, 1))
-        
-        #energy = [batch size, trg len, src len]
-        
-        attention = F.softmax(energy, dim=2)
-        
-        #attention = [batch size, trg len, src len]
-            
-        attended_encoding = torch.matmul(attention, encoder_combined)
-        
-        #attended_encoding = [batch size, trg len, emd dim]
-        
-        #convert from emb dim -> hid dim
-        attended_encoding = self.attn_emb2hid(attended_encoding)
-        
-        #attended_encoding = [batch size, trg len, hid dim]
-        
-        #apply residual connection
-        attended_combined = (conved + attended_encoding.permute(0, 2, 1)) * self.scale
-        
-        #attended_combined = [batch size, hid dim, trg len]
-        
-        return attention, attended_combined
-        
-    def forward(self, trg, encoder_conved, encoder_combined):
-        
-        #trg = [batch size, trg len]
-        #encoder_conved = encoder_combined = [batch size, src len, emb dim]
-        trg = trg.transpose(0, 1)
-        
-        batch_size = trg.shape[0]
-        trg_len = trg.shape[1]
-        device = trg.device
-        
-        #create position tensor
-        pos = torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size, 1).to(device)
-        
-        #pos = [batch size, trg len]
-        
-        #embed tokens and positions
-        tok_embedded = self.tok_embedding(trg)
-        pos_embedded = self.pos_embedding(pos)
-        
-        #tok_embedded = [batch size, trg len, emb dim]
-        #pos_embedded = [batch size, trg len, emb dim]
-        
-        #combine embeddings by elementwise summing
-        embedded = self.dropout(tok_embedded + pos_embedded)
-        
-        #embedded = [batch size, trg len, emb dim]
-        
-        #pass embedded through linear layer to go through emb dim -> hid dim
-        conv_input = self.emb2hid(embedded)
-        
-        #conv_input = [batch size, trg len, hid dim]
-        
-        #permute for convolutional layer
-        conv_input = conv_input.permute(0, 2, 1) 
-        
-        #conv_input = [batch size, hid dim, trg len]
-        
-        batch_size = conv_input.shape[0]
-        hid_dim = conv_input.shape[1]
-        
-        for i, conv in enumerate(self.convs):
-        
-            #apply dropout
-            conv_input = self.dropout(conv_input)
-        
-            #need to pad so decoder can't "cheat"
-            padding = torch.zeros(batch_size, 
-                                  hid_dim, 
-                                  self.kernel_size - 1).fill_(self.trg_pad_idx).to(device)
-                
-            padded_conv_input = torch.cat((padding, conv_input), dim = 2)
-        
-            #padded_conv_input = [batch size, hid dim, trg len + kernel size - 1]
-        
-            #pass through convolutional layer
-            conved = conv(padded_conv_input)
-
-            #conved = [batch size, 2 * hid dim, trg len]
-            
-            #pass through GLU activation function
-            conved = F.glu(conved, dim = 1)
-
-            #conved = [batch size, hid dim, trg len]
-            
-            #calculate attention
-            attention, conved = self.calculate_attention(embedded, 
-                                                         conved, 
-                                                         encoder_conved, 
-                                                         encoder_combined)
-            
-            #attention = [batch size, trg len, src len]
-            
-            #apply residual connection
-            conved = (conved + conv_input) * self.scale
-            
-            #conved = [batch size, hid dim, trg len]
-            
-            #set conv_input to conved for next loop iteration
-            conv_input = conved
-            
-        conved = self.hid2emb(conved.permute(0, 2, 1))
-         
-        #conved = [batch size, trg len, emb dim]
-            
-        output = self.fc_out(self.dropout(conved))
-        
-        #output = [batch size, trg len, output dim]
-            
-        return output, attention
-
-class ConvSeq2Seq(nn.Module):
-    def __init__(self, vocab_size, emb_dim, hid_dim, enc_layers, dec_layers, enc_kernel_size, dec_kernel_size, enc_max_length, dec_max_length, dropout, pad_idx, device):
-        super().__init__()
-        
-        enc = Encoder(emb_dim, hid_dim, enc_layers, enc_kernel_size, dropout, device, enc_max_length)
-        dec = Decoder(vocab_size, emb_dim, hid_dim, dec_layers, dec_kernel_size, dropout, pad_idx, device, dec_max_length)
-
-        self.encoder = enc
-        self.decoder = dec
-
-    def forward_encoder(self, src):
-        encoder_conved, encoder_combined = self.encoder(src)
-        
-        return encoder_conved, encoder_combined
-
-    def forward_decoder(self, trg, memory):
-        encoder_conved, encoder_combined = memory
-        output, attention = self.decoder(trg, encoder_conved, encoder_combined)
-        
-        return output, (encoder_conved, encoder_combined) 
-
-    def forward(self, src, trg):
-        
-        #src = [batch size, src len]
-        #trg = [batch size, trg len - 1] (<eos> token sliced off the end)
-           
-        #calculate z^u (encoder_conved) and (z^u + e) (encoder_combined)
-        #encoder_conved is output from final encoder conv. block
-        #encoder_combined is encoder_conved plus (elementwise) src embedding plus 
-        #  positional embeddings 
-        encoder_conved, encoder_combined = self.encoder(src)
-            
-        #encoder_conved = [batch size, src len, emb dim]
-        #encoder_combined = [batch size, src len, emb dim]
-        
-        #calculate predictions of next words
-        #output is a batch of predictions for each word in the trg sentence
-        #attention a batch of attention scores across the src sentence for 
-        #  each word in the trg sentence
-        output, attention = self.decoder(trg, encoder_conved, encoder_combined)
-        
-        #output = [batch size, trg len - 1, output dim]
-        #attention = [batch size, trg len - 1, src len]
-        
-        return output#, attention

vietocr/model/seqmodel/seq2seq.py

@@ -1,175 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torch.nn.functional as F
-
-class Encoder(nn.Module):
-    def __init__(self, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
-        super().__init__()
-                
-        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional = True)
-        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
-        self.dropout = nn.Dropout(dropout)
-        
-    def forward(self, src):
-        """
-        src: src_len x batch_size x img_channel
-        outputs: src_len x batch_size x hid_dim 
-        hidden: batch_size x hid_dim
-        """
-
-        embedded = self.dropout(src)
-        
-        outputs, hidden = self.rnn(embedded)
-                                 
-        hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1)))
-        
-        return outputs, hidden
-
-class Attention(nn.Module):
-    def __init__(self, enc_hid_dim, dec_hid_dim):
-        super().__init__()
-        
-        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)
-        self.v = nn.Linear(dec_hid_dim, 1, bias = False)
-        
-    def forward(self, hidden, encoder_outputs):
-        """
-        hidden: batch_size x hid_dim
-        encoder_outputs: src_len x batch_size x hid_dim,
-        outputs: batch_size x src_len
-        """
-        
-        batch_size = encoder_outputs.shape[1]
-        src_len = encoder_outputs.shape[0]
-        
-        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
-  
-        encoder_outputs = encoder_outputs.permute(1, 0, 2)
-        
-        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim = 2))) 
-        
-        attention = self.v(energy).squeeze(2)
-        
-        return F.softmax(attention, dim = 1)
-
-class Decoder(nn.Module):
-    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
-        super().__init__()
-
-        self.output_dim = output_dim
-        self.attention = attention
-        
-        self.embedding = nn.Embedding(output_dim, emb_dim)
-        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
-        self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)
-        self.dropout = nn.Dropout(dropout)
-        
-    def forward(self, input, hidden, encoder_outputs):
-        """
-        inputs: batch_size
-        hidden: batch_size x hid_dim
-        encoder_outputs: src_len x batch_size x hid_dim
-        """
-             
-        input = input.unsqueeze(0)
-        
-        embedded = self.dropout(self.embedding(input))
-        
-        a = self.attention(hidden, encoder_outputs)
-                
-        a = a.unsqueeze(1)
-        
-        encoder_outputs = encoder_outputs.permute(1, 0, 2)
-        
-        weighted = torch.bmm(a, encoder_outputs)
-        
-        weighted = weighted.permute(1, 0, 2)
-        
-        rnn_input = torch.cat((embedded, weighted), dim = 2)
-        
-        output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
-        
-        assert (output == hidden).all()
-        
-        embedded = embedded.squeeze(0)
-        output = output.squeeze(0)
-        weighted = weighted.squeeze(0)
-        
-        prediction = self.fc_out(torch.cat((output, weighted, embedded), dim = 1))
-        
-        return prediction, hidden.squeeze(0), a.squeeze(1)
-
-class Seq2Seq(nn.Module):
-    def __init__(self, vocab_size, encoder_hidden, decoder_hidden, img_channel, decoder_embedded, dropout=0.1):
-        super().__init__()
-        
-        attn = Attention(encoder_hidden, decoder_hidden)
-        
-        self.encoder = Encoder(img_channel, encoder_hidden, decoder_hidden, dropout)
-        self.decoder = Decoder(vocab_size, decoder_embedded, encoder_hidden, decoder_hidden, dropout, attn)
-        
-    def forward_encoder(self, src):       
-        """
-        src: timestep x batch_size x channel
-        hidden: batch_size x hid_dim
-        encoder_outputs: src_len x batch_size x hid_dim
-        """
-
-        encoder_outputs, hidden = self.encoder(src)
-
-        return (hidden, encoder_outputs)
-
-    def forward_decoder(self, tgt, memory):
-        """
-        tgt: timestep x batch_size 
-        hidden: batch_size x hid_dim
-        encouder: src_len x batch_size x hid_dim
-        output: batch_size x 1 x vocab_size
-        """
-        
-        tgt = tgt[-1]
-        hidden, encoder_outputs = memory
-        output, hidden, _ = self.decoder(tgt, hidden, encoder_outputs)
-        output = output.unsqueeze(1)
-        
-        return output, (hidden, encoder_outputs)
-
-    def forward(self, src, trg):
-        """
-        src: time_step x batch_size
-        trg: time_step x batch_size
-        outputs: batch_size x time_step x vocab_size
-        """
-
-        batch_size = src.shape[1]
-        trg_len = trg.shape[0]
-        trg_vocab_size = self.decoder.output_dim
-        device = src.device
-
-        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(device)
-        encoder_outputs, hidden = self.encoder(src)
-                
-        for t in range(trg_len):
-            input = trg[t] 
-            output, hidden, _ = self.decoder(input, hidden, encoder_outputs)
-            
-            outputs[t] = output
-            
-        outputs = outputs.transpose(0, 1).contiguous()
-
-        return outputs
-
-    def expand_memory(self, memory, beam_size):
-        hidden, encoder_outputs = memory
-        hidden = hidden.repeat(beam_size, 1)
-        encoder_outputs = encoder_outputs.repeat(1, beam_size, 1)
-
-        return (hidden, encoder_outputs)
-    
-    def get_memory(self, memory, i):
-        hidden, encoder_outputs = memory
-        hidden = hidden[[i]]
-        encoder_outputs = encoder_outputs[:, [i],:]
-
-        return (hidden, encoder_outputs)

vietocr/model/seqmodel/transformer.py

@@ -1,124 +0,0 @@
-from einops import rearrange
-from torchvision import models
-import math
-import torch
-from torch import nn
-
-class LanguageTransformer(nn.Module):
-    def __init__(self, vocab_size, 
-                 d_model, nhead, 
-                 num_encoder_layers, num_decoder_layers, 
-                 dim_feedforward, max_seq_length, 
-                 pos_dropout, trans_dropout):
-        super().__init__()
-        
-        self.d_model = d_model
-        self.embed_tgt = nn.Embedding(vocab_size, d_model)
-        self.pos_enc = PositionalEncoding(d_model, pos_dropout, max_seq_length)
-#        self.learned_pos_enc = LearnedPositionalEncoding(d_model, pos_dropout, max_seq_length)
-
-        self.transformer = nn.Transformer(d_model, nhead, 
-                                          num_encoder_layers, num_decoder_layers, 
-                                          dim_feedforward, trans_dropout)
-        
-        self.fc = nn.Linear(d_model, vocab_size)
-        
-    def forward(self, src, tgt, src_key_padding_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None):
-        """
-        Shape:
-            - src: (W, N, C)
-            - tgt: (T, N) 
-            - src_key_padding_mask: (N, S)
-            - tgt_key_padding_mask: (N, T)
-            - memory_key_padding_mask: (N, S)
-            - output: (N, T, E)
-            
-        """
-        tgt_mask = self.gen_nopeek_mask(tgt.shape[0]).to(src.device)
-        
-        src = self.pos_enc(src*math.sqrt(self.d_model))
-#        src = self.learned_pos_enc(src*math.sqrt(self.d_model))
-
-        tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))
-        
-        output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,
-                                  tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)
-#        output = rearrange(output, 't n e -> n t e')
-        output = output.transpose(0, 1)
-        return self.fc(output)
-
-    def gen_nopeek_mask(self, length):
-        mask = (torch.triu(torch.ones(length, length)) == 1).transpose(0, 1)
-        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
-
-        return mask
-    
-    def forward_encoder(self, src):
-        src = self.pos_enc(src*math.sqrt(self.d_model))
-        memory = self.transformer.encoder(src)
-        return memory
-    
-    def forward_decoder(self, tgt, memory):
-        tgt_mask = self.gen_nopeek_mask(tgt.shape[0]).to(tgt.device)
-        tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))
-        
-        output = self.transformer.decoder(tgt, memory, tgt_mask=tgt_mask)
-#        output = rearrange(output, 't n e -> n t e')
-        output = output.transpose(0, 1)
-
-        return self.fc(output), memory
-    
-    def expand_memory(self, memory, beam_size):
-        memory = memory.repeat(1, beam_size, 1)
-        return memory
-    
-    def get_memory(self, memory, i):
-        memory = memory[:, [i], :]
-        return memory
-
-class PositionalEncoding(nn.Module):
-    def __init__(self, d_model, dropout=0.1, max_len=100):
-        super(PositionalEncoding, self).__init__()
-        self.dropout = nn.Dropout(p=dropout)
-
-        pe = torch.zeros(max_len, d_model)
-        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0).transpose(0, 1)
-        self.register_buffer('pe', pe)
-
-    def forward(self, x):
-        x = x + self.pe[:x.size(0), :]
-
-        return self.dropout(x)
- 
-class LearnedPositionalEncoding(nn.Module):
-    def __init__(self, d_model, dropout=0.1, max_len=100):
-        super(LearnedPositionalEncoding, self).__init__()
-        self.dropout = nn.Dropout(p=dropout)
-
-        self.pos_embed = nn.Embedding(max_len, d_model)
-        self.layernorm = LayerNorm(d_model)
-
-    def forward(self, x):
-        seq_len = x.size(0)
-        pos = torch.arange(seq_len, dtype=torch.long, device=x.device)
-        pos = pos.unsqueeze(-1).expand(x.size()[:2])
-        x = x + self.pos_embed(pos)
-        return self.dropout(self.layernorm(x))
-
-class LayerNorm(nn.Module):
-    "A layernorm module in the TF style (epsilon inside the square root)."
-    def __init__(self, d_model, variance_epsilon=1e-12):
-        super().__init__()
-        self.gamma = nn.Parameter(torch.ones(d_model))
-        self.beta  = nn.Parameter(torch.zeros(d_model))
-        self.variance_epsilon = variance_epsilon
-
-    def forward(self, x):
-        u = x.mean(-1, keepdim=True)
-        s = (x - u).pow(2).mean(-1, keepdim=True)
-        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
-        return self.gamma * x + self.beta  

vietocr/model/transformerocr.py

@@ -1,44 +0,0 @@
-from vietocr.model.backbone.cnn import CNN
-from vietocr.model.seqmodel.transformer import LanguageTransformer
-from vietocr.model.seqmodel.seq2seq import Seq2Seq
-from vietocr.model.seqmodel.convseq2seq import ConvSeq2Seq
-from torch import nn
-
-class VietOCR(nn.Module):
-    def __init__(self, vocab_size,
-                 backbone,
-                 cnn_args, 
-                 transformer_args, seq_modeling='transformer'):
-        
-        super(VietOCR, self).__init__()
-        
-        self.cnn = CNN(backbone, **cnn_args)
-        self.seq_modeling = seq_modeling
-
-        if seq_modeling == 'transformer':
-            self.transformer = LanguageTransformer(vocab_size, **transformer_args)
-        elif seq_modeling == 'seq2seq':
-            self.transformer = Seq2Seq(vocab_size, **transformer_args)
-        elif seq_modeling == 'convseq2seq':
-            self.transformer = ConvSeq2Seq(vocab_size, **transformer_args)
-        else:
-            raise('Not Support Seq Model')
-
-    def forward(self, img, tgt_input, tgt_key_padding_mask):
-        """
-        Shape:
-            - img: (N, C, H, W)
-            - tgt_input: (T, N)
-            - tgt_key_padding_mask: (N, T)
-            - output: b t v
-        """
-        src = self.cnn(img)
-
-        if self.seq_modeling == 'transformer':
-            outputs = self.transformer(src, tgt_input, tgt_key_padding_mask=tgt_key_padding_mask)
-        elif self.seq_modeling == 'seq2seq':
-            outputs = self.transformer(src, tgt_input)
-        elif self.seq_modeling == 'convseq2seq': 
-            outputs = self.transformer(src, tgt_input)
-        return outputs
-

vietocr/model/vocab.py

@@ -1,36 +0,0 @@
-class Vocab():
-    def __init__(self, chars):
-        self.pad = 0
-        self.go = 1
-        self.eos = 2
-        self.mask_token = 3
-
-        self.chars = chars
-
-        self.c2i = {c:i+4 for i, c in enumerate(chars)}
-
-        self.i2c = {i+4:c for i, c in enumerate(chars)}
-        
-        self.i2c[0] = '<pad>'
-        self.i2c[1] = '<sos>'
-        self.i2c[2] = '<eos>'
-        self.i2c[3] = '*'
-
-    def encode(self, chars):
-        return [self.go] + [self.c2i[c] for c in chars] + [self.eos]
-    
-    def decode(self, ids):
-        first = 1 if self.go in ids else 0
-        last = ids.index(self.eos) if self.eos in ids else None
-        sent = ''.join([self.i2c[i] for i in ids[first:last]])
-        return sent
-    
-    def __len__(self):
-        return len(self.c2i) + 4
-    
-    def batch_decode(self, arr):
-        texts = [self.decode(ids) for ids in arr]
-        return texts
-
-    def __str__(self):
-        return self.chars