Добавил комментарии к некоторым функциям и классам

Author: Iliaavilov

dataset.py

@@ -7,32 +7,28 @@
 
 
 class CaptionDataset(Dataset):
-    """
-    A PyTorch Dataset class to be used in a PyTorch DataLoader to create batches.
-    """
 
-    def __init__(self, data_folder, split, transform=None):
-        """
-        :param data_folder: folder where data files are stored
-        :param split: split, one of 'TRAIN', 'VAL', or 'TEST'
-        :param transform: image transform pipeline
-        """
+    def __init__(self, data_folder, split):
+        '''
+        :param data_folder: Folder where files are stored
+        :param split: which split of dataset (train, validation or test)
+        '''
         self.split = split
         assert self.split in {'TRAIN', 'VAL', 'TEST'}
 
         # Open hdf5 file where images are stored
-        self.h = h5py.File(os.path.join(data_folder, self.split + '_IMAGES_' + 'coco.hdf5'), 'r')
+        self.h = h5py.File(os.path.join(data_folder, self.split + '_IMAGES_' + 'flickr30k.hdf5'), 'r')
         self.imgs = self.h['images']
 
         # Captions per image
         self.cpi = self.h.attrs['captions_per_image']
 
         # Load encoded captions (completely into memory)
-        with open(os.path.join(data_folder, self.split + '_CAPTIONS_' + 'coco.json'), 'r') as j:
+        with open(os.path.join(data_folder, self.split + '_CAPTIONS_' + 'flickr30k.json'), 'r') as j:
             self.captions = json.load(j)
 
         # Load caption lengths (completely into memory)
-        with open(os.path.join(data_folder, self.split + '_CAPLENS_' + 'coco.json'), 'r') as j:
+        with open(os.path.join(data_folder, self.split + '_CAPLENS_' + 'flickr30k.json'), 'r') as j:
             self.caplens = json.load(j)
 
         # PyTorch transformation pipeline for the image (normalizing, etc.)
@@ -61,4 +57,5 @@ def __getitem__(self, i):
             return img, caption, caplen, all_captions
 
     def __len__(self):
+
         return self.dataset_size

model.py

@@ -7,26 +7,30 @@
 class Encoder(nn.Module):
 
     def __init__(self, encoded_image_size=14):
+        '''
+        :param encoded_image_size: each encoded channel size of image will be encoded_image_size X encoded_image_size
+        '''
+
         super(Encoder, self).__init__()
         self.enc_image_size = encoded_image_size
-
-        # Загружаем натренированную resnet152
+        # Load pretrained resnet model
         resnet = torchvision.models.resnet152(pretrained=True)
-        # Убираем линейные слои (нам нужны только CNN)
+        # Delete FC layers and leave only CNN
         modules = list(resnet.children())[:-2]
         self.resnet = nn.Sequential(*modules)
-        # Ресайз фичей изображения к нужным размерам
+        # Resize CNN features from resnet to appropriate size
         self.pooling = nn.AdaptiveAvgPool2d((encoded_image_size, encoded_image_size))
+        # Resnet parameters will not be modified during training
         for p in self.resnet.parameters():
             p.requires_grad = False
 
     def forward(self, images):
 
-        # Извлекаем 2048 "каналов" фичей по 7X7 каждый
+        # Obtain 2048 channels of features each of size 7X7
         out = self.resnet(images)  # (batch_size, 2048, 7, 7)
-        # Изменяем размер каналов до (encoded_image_size, encoded_image_size)
+        # Reseize size to (encoded_image_size, encoded_image_size)
         out = self.pooling(out)  # (batch_size, 2048, encoded_image_size, encoded_image_size)
-        # Переставляем местами размерности (просто для удобства)
+        # Change dimension places (just for convinience)
         out = out.permute(0, 2, 3, 1)  # (batch_size, encoded_image_size, encoded_image_size, 2048)
         return out
 
@@ -36,19 +40,23 @@ class Attention(nn.Module):
     def __init__(self, word_embeddings_dim, attention_dim, encoded_image_size):
         super(Attention, self).__init__()
 
+        # Attention layer for encoder
         self.att_encoder = nn.Linear(2048, attention_dim)
+        # Attention layer for decoder
         self.att_decoder = torch.nn.Linear(word_embeddings_dim, attention_dim)
+        # Final layer of attention
         self.att_final = torch.nn.Linear(attention_dim, 1)
         self.softmax = nn.Softmax(dim = 1)
         self.relu = nn.ReLU()
 
-    def forward(self, encoder_out, decoder_out, batch_size):
+    def forward(self, encoder_out, decoder_out):
+
         att_encoder_computed = self.att_encoder(encoder_out)  # (batch_size, encoded_image_size**2, attention_dim)
         att_decoder_computed = self.att_decoder(decoder_out)  # (batch_size, attention_dim)
         att = self.att_final(self.relu(att_encoder_computed + att_decoder_computed.unsqueeze(1))).squeeze(2)  # (batch_size, encoded_image_size**2)
         att_weights = self.softmax(att)  # (batch_size, 2048)
-
         encoder_weighted = (encoder_out * att_weights.unsqueeze(2)).sum(dim=1)  # (batch_size, encoder_dim)
+
         return encoder_weighted
 
 
@@ -62,13 +70,11 @@ def __init__(self, vocab_size, word_embeddings_dim, attention_dim, decoder_hidde
         self.word_embeddings_dim = word_embeddings_dim
         self.vocab_size = vocab_size
         self.encoded_image_size = encoded_image_size
-
         self.LSTMCell = torch.nn.LSTMCell(2048 + word_embeddings_dim,
                                           hidden_size=decoder_hidden_size, bias = True)
         self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=word_embeddings_dim)
         self.Attention = Attention(word_embeddings_dim, attention_dim, encoded_image_size)
         self.linear = torch.nn.Linear(decoder_hidden_size, vocab_size)
-
         self.h_init = torch.nn.Linear(2048, decoder_hidden_size)
         self.c_init = torch.nn.Linear(2048, decoder_hidden_size)
         self.f_beta = nn.Linear(decoder_hidden_size, 2048)  # linear layer to create a sigmoid-activated gate

model_training.py

@@ -3,88 +3,119 @@
 from nltk.translate.bleu_score import corpus_bleu
 from utils import save_checkpoint
 import neptune
-import time
 
-def validate(enc, dec, device, loss_fn, val_loader, wordmap, epoch):
+def train(enc, dec, device, loss_fn, train_loader, optimizer_decoder, optimizer_encoder, epoch):
+    ''' Train model
+    :param enc: encoder part of model
+    :param dec: decoder part of model
+    :param device: on which device to train model
+    :param loss_fn: loss function
+    :param train_loader: pytorch loader of images
+    :param optimizer_decoder: pytorch optimizer for decoder part of model
+    :param optimizer_encoder: pytorch optimizer for encoder part of model
+    :param epoch: current epoch of training
+    :return: None
+    '''
+
+    dec.train()
+    enc.train()
+
+    dec = dec.to(device)
+    enc = enc.to(device)
+
+    # iterate through batches of train loader
+    for batch_n, (imgs, caps, caplens) in enumerate(train_loader):
+
+        imgs = imgs.to(device)
+        caps = caps.to(device)
+        caplens = caplens.to(device)
+
+        # Encode images
+        enc_output = enc(imgs)
+
+        # Decode encodings and get captions
+        dec_out, captions, captions_lengths, sort_ind = dec(captions=caps,
+                                                            encoder_out=enc_output,
+                                                            captions_lengths=caplens)
+
+        # Remove words which we did not decode at (e.g. max length of sentence in batch is 15 words,
+        # so for sentence of 10 words we did not decode 5 words, and we have to skip them during loss computing)
+        dec_out = pack_padded_sequence(dec_out, captions_lengths.cpu(), batch_first=True).data.to(device)
+        captions = pack_padded_sequence(captions, captions_lengths.cpu(), batch_first=True).data.to(device)
+
+        loss = loss_fn(dec_out, captions)
+        optimizer_decoder.zero_grad()
+        optimizer_encoder.zero_grad()
+
+        loss.backward()
+
+        optimizer_decoder.step()
+        optimizer_encoder.step()
+
+        if batch_n % 3000 == 0:
+            save_checkpoint(epoch, batch_n, enc, dec, optimizer_encoder, optimizer_decoder)
+            print('Current loss', loss.item())
+
+    # Log metric to neptune
+    neptune.log_metric('loss', loss.item())
+
+
+def validate(enc, dec, device, val_loader, wordmap, epoch):
+    ''' Calculate validation metric
+    :param val_loader: pytorch loader of images
+    :param wordmap: dictionary mapping from word to word index
+    :param epoch: current epoch of training
+    :return: None
+    '''
+
     enc.eval()
     dec.eval()
-    references = list()  # references (true captions) for calculating BLEU-4 score
-    hypotheses = list()  # hypotheses (predictions)
+
+    dec = dec.to(device)
+    enc = enc.to(device)
+
+    references = list()  # True captions
+    hypotheses = list()  # Predicted captions
+
     with torch.no_grad():
+
         for batch_n, (imgs, caps, caplens, allcaps) in enumerate(val_loader):
+
             print(batch_n)
+
             imgs = imgs.to(device)
             caps = caps.to(device)
             caplens = caplens.to(device)
+
             enc_output = enc(imgs)
             dec_out, captions, captions_lengths, sort_ind = dec(captions=caps,
                                                                 encoder_out=enc_output,
                                                                 captions_lengths=caplens)
-
             scores_copy = dec_out.clone()
-            dec_out = pack_padded_sequence(dec_out.cpu(), captions_lengths.cpu(), batch_first=True).data.to(device)
-            captions = pack_padded_sequence(captions.cpu(), captions_lengths.cpu(), batch_first=True).data.to(device)
 
-            loss = loss_fn(dec_out, captions)
 
-            allcaps = allcaps[sort_ind]  # because images were sorted in the decoder
+            allcaps = allcaps[sort_ind]  # Resort because captions were sorted in decoder
+
             for j in range(allcaps.shape[0]):
+
                 img_caps = allcaps[j].tolist()
                 img_captions = list(
                     map(lambda c: [w for w in c if w not in {wordmap['<start>'], wordmap['<pad>']}],
                         img_caps))  # remove <start> and pads
                 references.append(img_captions)
 
-            # Hypotheses
+            # Take predicted captions for each image
             _, preds = torch.max(scores_copy, dim=2)
             preds = preds.tolist()
             temp_preds = list()
             for j, p in enumerate(preds):
                 temp_preds.append(preds[j][:captions_lengths[j]])  # remove pads
             preds = temp_preds
             hypotheses.extend(preds)
+
         # Calculate BLEU-4 scores
         bleu4 = corpus_bleu(references, hypotheses)
-        neptune.log_metric('bleu4', bleu4)
-        print('Epoch {}, BLEU4'.format(epoch), bleu4)
-
-
-def train(enc, dec, device, loss_fn, train_loader, optimizer_decoder, optimizer_encoder, epoch):
-
-    dec.train()  # train mode (dropout and batchnorm is used)
-    enc.train()
-
-    for batch_n, (imgs, caps, caplens) in enumerate(train_loader):
-        start = time.time()
-        imgs = imgs.to(device)
-        caps = caps.to(device)
-        caplens = caplens.to(device)
-        enc_output = enc(imgs)
-        dec_out, captions, captions_lengths, sort_ind = dec(captions=caps,
-                                                            encoder_out=enc_output,
-                                                            captions_lengths=caplens)
-        #        if batch_n % 20 == 0:
-        #            aaaa = [res.get(int(key)) for key in torch.argmax(dec_out[0], dim = 1)]
-        #            print('epoch:', epoch, 'batch', batch_n, aaaa)
-        #            img = Image.fromarray((unorm(imgs[0].cpu()).numpy()*255).astype('uint8').transpose(1, 2, 0))
-        #            img.save('{}-{}.png'.format(epoch, batch_n))
-        #            with open("captions.txt", "a") as f:
-        #                # Append 'hello' at the end of file
-        #                f.write("\n")
-        #                f.write(str(epoch) + '_' + str(batch_n) + '_' + str(aaaa))
-        dec_out = pack_padded_sequence(dec_out, captions_lengths.cpu(), batch_first=True).data.to(device)
-        captions = pack_padded_sequence(captions, captions_lengths.cpu(), batch_first=True).data.to(device)
-
-        loss = loss_fn(dec_out, captions)
-        optimizer_decoder.zero_grad()
-        optimizer_encoder.zero_grad()
 
-        loss.backward()
-
-        optimizer_decoder.step()
-        optimizer_encoder.step()
-        if batch_n % 3000 == 0:
-            save_checkpoint(epoch, batch_n, enc, dec, optimizer_encoder, optimizer_decoder)
-            print('Current loss', loss.item())
-
-    neptune.log_metric('loss', loss.item())
+        # Log score to neptune and print metric
+        neptune.log_metric('bleu4', bleu4)
+        print('Epoch {}, BLEU4'.format(epoch), bleu4)

server/server.py

@@ -9,7 +9,7 @@
 
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-enc, dec = load_models(checkpoint_name = 'checkpoint_7_0.pth.tar')
+enc, dec = load_models(checkpoint_name = 'checkpoint_6_0.pth.tar')
 wordmap, res = load_wordmap()
 app = Flask(__name__)
 
@@ -23,7 +23,6 @@ def get_image_caption():
     no_image = False
     try:
         image = request.files['image']
-        print(type(image))
         image_hash = hash(image)
         image.save('images/{}.png'.format(image_hash))
         image = imread('images/{}.png'.format(image_hash))
@@ -34,8 +33,9 @@ def get_image_caption():
         image = image_preprocessing(image)
         if (image.shape == (3, 256, 256)) & (np.max(image) <= 256):
             image = image_normalisation(image, device)
-            predicted_captions, _ = captioning(enc, dec, image, wordmap, device, res)
-            return ' '.join(predicted_captions[1:])
+            predicted_captions, encoders_out = captioning(enc, dec, image, wordmap, device, res)
+            return jsonify({'captions': ' '.join(predicted_captions[1:]),
+                            'encoders_out': encoders_out.tolist()})
         else:
             return 'Image is not in png'
 
@@ -45,14 +45,13 @@ def get_video_captions():
     no_video = False
     try:
         video = request.files['video']
-        print(type(video))
         video_hash = hash(video)
         video.save('videos/{}.mp4'.format(video_hash))
     except:
         no_video = True
         return jsonify('No video file in post requests')
     if no_video == False:
-        video_to_screenshots('videos/{}.mp4'.format(video_hash), 'saved_screenshots', 200)
+        video_to_screenshots('videos/{}.mp4'.format(video_hash), 'saved_screenshots', 10)
         list_of_files = [f for f in listdir('saved_screenshots') if isfile(join('saved_screenshots', f))]
         all_captions = []
         all_encoders_out = []
@@ -63,25 +62,19 @@ def get_video_captions():
             if (image.shape == (3, 256, 256)) & (np.max(image) <= 256):
                 image = image_normalisation(image, device)
                 predicted_captions, encoder_out = captioning(enc, dec, image, wordmap, device, res)
-                all_captions.append(predicted_captions)
-                all_encoders_out.append(encoder_out)
+                all_captions.append(' '.join(predicted_captions[1:]))
+                all_encoders_out.append(encoder_out.tolist())
 
         return jsonify({'captions': all_captions,
                         'encoders_out': all_encoders_out})
 
-
-
-
-
-
 @app.route('/get_captions', methods=['POST', 'GET'])
 def get_captions():
 
     if request.method == 'POST':
         no_image = False
         try:
             image = request.files['image']
-            print(type(image))
             image_hash = hash(image)
             image.save('images/{}.png'.format(image_hash))
             image = imread('images/{}.png'.format(image_hash))