Add application file

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,19 @@
+FROM python:3.9
+
+# Buat user non-root
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+# Set direktori kerja
+WORKDIR /app
+
+# Salin dependencies dan install
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+# Salin semua file app ke image
+COPY --chown=user . /app
+
+# Jalankan aplikasi Streamlit
+CMD ["streamlit", "run", "app_streamlit.py", "--server.port=7860", "--server.address=0.0.0.0"]

app_streamlit.py

@@ -0,0 +1,173 @@
+import streamlit as st
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForVisualQuestionAnswering
+from PIL import Image, ImageOps
+import io
+import torchvision.transforms as T
+import torch.nn.functional as F
+
+# **Cek Device**
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# **Konfigurasi Halaman Streamlit**
+st.set_page_config(
+    initial_sidebar_state="expanded",
+    page_title="Explainable Image Caption Bot"
+)
+
+# **Load Model BLIP**
+@st.cache_resource
+def load_blip_model():
+    # processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+    # model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
+    processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
+    model = AutoModelForVisualQuestionAnswering.from_pretrained("Salesforce/blip2-opt-2.7b").to(device)
+    return processor, model
+
+processor, model = load_blip_model()
+
+# **Transformasi Gambar untuk Model**
+def transform_image(img):
+    transform = T.Compose([
+        T.Resize((384, 384)),  # Resize sesuai model BLIP
+        T.ToTensor(),
+        T.Normalize((0.5,), (0.5,))
+    ])
+    return transform(img)
+
+def generate_caption(image, processor, model):
+    inputs = processor(images=image, return_tensors="pt").to(device)
+
+    # Pastikan kita menangkap perhatian dari Transformer
+    attention_maps = []
+
+    def get_attention_hook(module, input, output):
+        print("✅ Hook executed! Attention captured.")  # Debugging
+        attention_maps.append(output)  # Output adalah tuple
+
+    # Pasang hook ke layer yang sesuai
+    handle = model.vision_model.encoder.layers[-1].self_attn.register_forward_hook(get_attention_hook)
+
+    # Generate caption
+    with torch.no_grad():
+        caption_ids = model.generate(**inputs)
+
+    # Hapus hook setelah digunakan
+    handle.remove()
+
+    caption = processor.decode(caption_ids[0], skip_special_tokens=True)
+
+    # **Periksa apakah attention_maps berhasil ditangkap**
+    if not attention_maps:
+        print("❌ Attention Maps tidak terisi! Hook mungkin tidak bekerja.")
+        return caption, None
+
+    # **Ambil tensor dari tuple**
+    attention_tensor = attention_maps[0][0]  # Ambil tensor pertama dari tuple
+    attention = attention_tensor.cpu().detach().numpy().mean(axis=1)
+
+    return caption, attention
+
+
+# **Fungsi untuk Memuat Gambar**
+@st.cache_data
+def load_uploaded_image(img):
+    if isinstance(img, str):
+        image = Image.open(img)
+    else:
+        img_bytes = img.read()
+        image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+    
+    image = ImageOps.exif_transpose(image)  # Perbaiki orientasi gambar
+    return image
+
+def plot_attention(image, caption, attention):
+    """
+    Menampilkan heatmap attention untuk setiap kata dalam caption.
+    """
+
+    if attention is None or len(attention.shape) != 2:
+        st.error("Attention map tidak valid! Tidak bisa menampilkan heatmap.")
+        return
+
+    num_words = len(caption.split())
+    num_attention_steps = min(num_words, attention.shape[0])  # Sesuaikan panjang attention
+
+    fig, axes = plt.subplots(1, num_attention_steps, figsize=(num_attention_steps * 3, 5))
+
+    if num_attention_steps == 1:
+        axes = [axes]  # Pastikan list jika hanya ada satu kata
+
+    for i in range(num_attention_steps):
+        attn_map = attention[i]
+
+        # **Reshape attention ke bentuk yang sesuai**
+        if attn_map.shape[0] == 768:
+            grid_size = 24  # Vision Transformer biasanya menggunakan 24x32 patches
+            attn_map = attn_map[:grid_size * grid_size].reshape(grid_size, grid_size)
+        else:
+            st.warning(f"Attention map tidak bisa diubah menjadi grid! (Token count: {attn_map.shape[0]})")
+            continue
+
+        # **Interpolasi agar ukuran sesuai dengan gambar**
+        attn_resized = F.interpolate(
+            torch.tensor(attn_map).unsqueeze(0).unsqueeze(0), 
+            size=(image.size[1], image.size[0]),  # Sesuaikan ke ukuran gambar
+            mode="bilinear",
+            align_corners=False
+        ).squeeze().numpy()
+
+        # **Plot setiap heatmap per kata**
+        axes[i].imshow(image)
+        axes[i].imshow(attn_resized, cmap='jet', alpha=0.5)
+        axes[i].set_title(caption.split()[i])
+        axes[i].axis("off")
+
+    plt.tight_layout()
+    st.pyplot(fig)
+
+# **Streamlit UI**
+st.title("Explainable Image Captioning Bot 🤖🖼️")
+st.text("Powered by BLIP (Salesforce) - A Transformer-based Image Captioning Model")
+
+st.success("Upload an image and generate a caption!")
+
+# **File Upload**
+uploaded_file = st.file_uploader("Upload an image (JPG, PNG, JPEG)", type=["png", "jpg", "jpeg", "webp"])
+img_path = "imgs/test2.jpeg" if uploaded_file is None else uploaded_file
+
+# **Muat dan Tampilkan Gambar**
+image = load_uploaded_image(img_path)
+st.image(image, use_column_width=True, caption="Uploaded Image")
+
+# **Generate Caption Button**
+# Jika tombol ditekan, jalankan captioning dan attention visualization
+if st.button("Generate Caption"):
+    caption, attention = generate_caption(image, processor, model)
+
+    if attention is None:
+        st.error("Attention map tidak tersedia! Coba ganti layer yang di-hook.")
+    else:
+        st.markdown(f"### **Generated Caption:**\n📢 *{caption}*")
+        plot_attention(image, caption, attention)  # ✅ Panggil dengan 3 argumen
+
+    st.balloons()
+
+
+# **Sidebar Info**
+st.sidebar.markdown("""
+### About This App 📝
+This app generates captions for images using **Hugging Face's BLIP model** trained by **Salesforce**.  
+It also provides **explainable AI insights** into how images are understood by deep learning models.
+
+### How to Use:
+1. **Upload an image** 📷 (JPG/PNG/JPEG).
+2. **Click "Generate Caption"** 🏷️.
+3. **View AI-generated caption** for your image along with **attention heatmap**!
+
+### Want More Features?
+Check the model on [Hugging Face](https://huggingface.co/Salesforce/blip-image-captioning-base).
+""")

attention_model_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83319fd76020f7b4a562b0b607c8b89fc8a604b5d385d72ea38c4b1b9c36d5b4
+size 230166330

caption.py

@@ -0,0 +1,120 @@
+from nltk.translate.bleu_score import corpus_bleu
+from tqdm import tqdm
+from dataset import Vocabulary
+from skimage import transform
+from model import *
+from utils import *
+import torchvision.transforms as T
+from PIL import Image
+import argparse
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+# Will only work for batch size 1
+def get_all_captions(img, model, vocab=None):
+    features = model.EncoderCNN(img[0:1].to(device))
+    caps, alphas = model.DecoderLSTM.gen_captions(features, vocab=vocab)
+    caps = caps[:-2]
+    return caps
+
+
+def calculate_bleu_score(dataloader, model, vocab):
+    candidate_corpus = []
+    references_corpus = []
+
+    for batch in tqdm(dataloader, total=len(dataloader)):
+        img, cap, all_caps = batch
+        img, cap = img.to(device), cap.to(device)
+        caps = get_all_captions(img, model, vocab)
+        candidate_corpus.append(caps)
+        references_corpus.append(all_caps[0])
+
+    assert len(candidate_corpus) == len(references_corpus)
+    print(f"\nBLEU1 = {corpus_bleu(references_corpus, candidate_corpus, (1, 0, 0, 0))}")
+    print(f"BLEU2 = {corpus_bleu(references_corpus, candidate_corpus, (0.5, 0.5, 0, 0))}")
+    print(f"BLEU3 = {corpus_bleu(references_corpus, candidate_corpus, (0.33, 0.33, 0.33, 0))}")
+    print(f"BLEU4 = {corpus_bleu(references_corpus, candidate_corpus, (0.25, 0.25, 0.25, 0.25))}")
+
+
+def get_caps_from(features_tensors, model, vocab=None):
+    model.eval()
+    with torch.no_grad():
+        features = model.EncoderCNN(features_tensors[0:1].to(device))
+        caps, alphas = model.DecoderLSTM.gen_captions(features, vocab=vocab)
+        caption = ' '.join(caps)
+        show_img(features_tensors[0], caption)
+
+    return caps, alphas
+
+
+def plot_attention(img, target, attention_plot):
+    img = img.to('cpu').numpy().transpose((1, 2, 0))
+    temp_image = img
+
+    fig = plt.figure(figsize=(15, 15))
+    len_caps = len(target)
+    for i in range(len_caps):
+        temp_att = attention_plot[i].reshape(7, 7)
+        temp_att = transform.pyramid_expand(temp_att, upscale=24, sigma=8)
+        ax = fig.add_subplot(len_caps // 2, len_caps // 2, i + 1)
+        ax.set_title(target[i])
+        img = ax.imshow(temp_image)
+        ax.imshow(temp_att, cmap='gray', alpha=0.5, extent=img.get_extent())
+
+    plt.tight_layout()
+    plt.show()
+
+
+def plot_caption_with_attention(img_pth, model, transforms_=None, vocab=None):
+    img = Image.open(img_pth)
+    img = transforms_(img)
+    img.unsqueeze_(0)
+    caps, attention = get_caps_from(img, model, vocab)
+    plot_attention(img[0], caps, attention)
+
+
+def main(arguments):
+    state_checkpoint = torch.load(arguments.state_chechpoint, map_location=device)  # change paths
+    # model params
+    vocab = state_checkpoint['vocab']
+    embed_size = arguments.embed_size
+    embed_wts = None
+    vocab_size = state_checkpoint['vocab_size']
+    attention_dim = arguments.attention_dim
+    encoder_dim = arguments.encoder_dim
+    decoder_dim = arguments.decoder_dim
+    fc_dims = arguments.fc_dims
+
+    model = EncoderDecoder(embed_size,
+                           vocab_size,
+                           attention_dim,
+                           encoder_dim,
+                           decoder_dim,
+                           fc_dims,
+                           p=0.3,
+                           embeddings=embed_wts).to(device)
+
+    model.load_state_dict(state_checkpoint['state_dict'])
+
+    transforms = T.Compose([
+        T.Resize((224, 224)),
+        T.ToTensor(),
+        T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+    ])
+
+    img_path = arguments.image
+    plot_caption_with_attention(img_path, model, transforms, vocab)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--image', type=str, required=True, help='input image for generating caption')
+    parser.add_argument('--state_checkpoint', type=str, required=True, help='path for state checkpoint')
+    parser.add_argument('--embed_size', type=int, default=300, help='dimension of word embedding vectors')
+    parser.add_argument('--attention_dim', type=int, default=256, help='dimension of attention layer')
+    parser.add_argument('--encoder_dim', type=int, default=2048, help='dimension of encoder layer')
+    parser.add_argument('--decoder_dim', type=int, default=512, help='dimension of decoder layer')
+    parser.add_argument('--fc_dims', type=int, default=256, help='dimension of fully connected layer')
+    args = parser.parse_args()
+    main(args)

dataset.py

@@ -0,0 +1,91 @@
+import torch
+import pandas as pd
+from torch.nn.utils.rnn import pad_sequence
+from PIL import Image
+import spacy
+import os
+
+from torch.utils.data import Dataset
+
+spacy_eng = spacy.load('en_core_web_sm')
+
+
+class Vocabulary:
+    def __init__(self, freq_threshold=5):
+        self.freq_threshold = freq_threshold
+        self.itos = {0: "<PAD>", 1: "<SOS>", 2: "<EOS>", 3: "<UNK>"}
+        self.stoi = {v: k for k, v in self.itos.items()}
+
+    def __len__(self):
+        return len(self.itos)
+
+    @staticmethod
+    def tokenize(text):
+        return [token.text.lower() for token in spacy_eng.tokenizer(text)]
+
+    def build_vocab(self, sent_list):
+        freqs = {}
+        idx = 4
+        for sent in sent_list:
+            for word in self.tokenize(sent):
+                if word not in freqs:
+                    freqs[word] = 1
+                else:
+                    freqs[word] += 1
+
+                if freqs[word] == self.freq_threshold:
+                    self.itos[idx] = word
+                    self.stoi[word] = idx
+                    idx += 1
+
+    def numericalize(self, sents):
+        tokens = self.tokenize(sents)
+        return [self.stoi[token] if token in self.stoi else self.stoi["<UNK>"]
+                for token in tokens]
+
+
+class FlickrDataset(Dataset):
+    def __init__(self, root_dir, csv_file, transforms=None, freq_threshold=5):
+        self.root_dir = root_dir
+        self.df = pd.read_csv(csv_file)
+        self.transforms = transforms
+
+        self.img_pts = self.df.iloc[:, 0]
+        self.caps = self.df.iloc[:, 1]
+        self.vocab = Vocabulary(freq_threshold)
+        self.vocab.build_vocab(self.caps.tolist())
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        captions = self.caps[idx]
+        img_pt = self.img_pts[idx]
+
+        img = Image.open(os.path.join(self.root_dir, img_pt)).convert('RGB')
+
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        encoded_cap = []
+        encoded_cap += [self.vocab.stoi["<SOS>"]]  # stoi string to index
+        encoded_cap += self.vocab.numericalize(captions)
+        encoded_cap += [self.vocab.stoi["<EOS>"]]
+        encoded_cap = torch.LongTensor(encoded_cap)
+
+        return img, encoded_cap
+
+
+class CapsCollate:
+    def __init__(self, pad_idx, batch_first=False):
+        self.pad_idx = pad_idx
+        self.batch_first = batch_first
+
+    def __call__(self, batch):
+        imgs = [item[0].unsqueeze(0) for item in batch]
+        imgs = torch.cat(imgs, dim=0)
+
+        targets = [item[1] for item in batch]
+        targets = pad_sequence(targets, batch_first=self.batch_first, padding_value=self.pad_idx)
+
+        return imgs, targets

download_files.py

@@ -0,0 +1,33 @@
+import requests
+
+
+def download_file_from_google_drive(id, destination):
+    URL = "https://docs.google.com/uc?export=download"
+
+    session = requests.Session()
+
+    response = session.get(URL, params={'id': id}, stream=True)
+    token = get_confirm_token(response)
+
+    if token:
+        params = {'id': id, 'confirm': token}
+        response = session.get(URL, params=params, stream=True)
+
+    save_response_content(response, destination)
+
+
+def get_confirm_token(response):
+    for key, value in response.cookies.items():
+        if key.startswith('download_warning'):
+            return value
+
+    return None
+
+
+def save_response_content(response, destination):
+    CHUNK_SIZE = 32768
+
+    with open(destination, "wb") as f:
+        for chunk in response.iter_content(CHUNK_SIZE):
+            if chunk:  # filter out keep-alive new chunks
+                f.write(chunk)

model.py

@@ -0,0 +1,134 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Encoder(nn.Module):
+    def __init__(self):
+        super(Encoder, self).__init__()
+        resnet = models.resnet50(pretrained=True)
+        for param in resnet.parameters():
+            param.requires_grad_(False)
+
+        modules = list(resnet.children())[:-2]  # extracting the last conv layer from the model
+        self.resnet = nn.Sequential(*modules)
+
+    def forward(self, imgs):
+        features = self.resnet(imgs)
+        features = features.permute(0, 2, 3, 1)  # batch x 7 x 7 x 2048
+        features = features.view(features.size(0), -1, features.size(-1))  # batch x 49 x 2048
+        return features
+
+
+class Attention(nn.Module):
+    def __init__(self, encoder_dims, decoder_dims, attention_dims):
+        super(Attention, self).__init__()
+        self.attention_dims = attention_dims  # size of attention network
+        self.U = nn.Linear(encoder_dims, attention_dims)  # a^(t)
+        self.W = nn.Linear(decoder_dims, attention_dims)  # s^(t` - 1)
+        self.A = nn.Linear(attention_dims, 1)  # cvt the attention dims back to 1
+
+    def forward(self, features, hidden):
+        u_as = self.U(features)
+        w_as = self.W(hidden)
+        combined_state = torch.tanh(u_as + w_as.unsqueeze(1))
+        attention_score = self.A(combined_state)
+        attention_score = attention_score.squeeze(2)
+        alpha = F.softmax(attention_score, dim=1)
+        attention_weights = features * alpha.unsqueeze(2)  # batch x num_timesteps (49) x features
+        attention_weights = attention_weights.sum(dim=1)
+        return alpha, attention_weights
+
+
+class Decoder(nn.Module):
+    def __init__(self, embed_size, vocab_size, attention_dim, encoder_dim, decoder_dim, fc_dims, p=0.3,
+                 embeddings=None):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.attention_dim = attention_dim
+        self.decoder_dim = decoder_dim
+
+        self.embedding = nn.Embedding(vocab_size, embedding_dim=embed_size)
+        self.attention = Attention(encoder_dim, decoder_dim, attention_dim)
+
+        self.init_h = nn.Linear(encoder_dim, decoder_dim)
+        self.init_c = nn.Linear(encoder_dim, decoder_dim)
+        self.lstm = nn.LSTMCell(encoder_dim + embed_size, decoder_dim, bias=True)
+        self.fcn1 = nn.Linear(decoder_dim, vocab_size)
+        self.fcn2 = nn.Linear(fc_dims, vocab_size)
+        self.drop = nn.Dropout(p)
+
+        if embeddings is not None:
+            self.load_pretrained_embed(embeddings)
+
+    def forward(self, features, captions):
+
+        seq_length = len(captions[0]) - 1  # Exclude the last one
+        batch_size = captions.size(0)
+        num_timesteps = features.size(1)
+
+        embed = self.embedding(captions)
+        h, c = self.init_hidden_state(features)  # initialize h and c for LSTM
+
+        preds = torch.zeros(batch_size, seq_length, self.vocab_size).to(device)
+        alphas = torch.zeros(batch_size, seq_length, num_timesteps).to(device)
+
+        for s in range(seq_length):
+            alpha, context = self.attention(features, h)
+            lstm_inp = torch.cat((embed[:, s], context), dim=1)
+            h, c = self.lstm(lstm_inp, (h, c))
+            out = self.drop(self.fcn1(h))
+            preds[:, s] = out
+            alphas[:, s] = alpha
+
+        return preds, alphas
+
+    def gen_captions(self, features, max_len=20, vocab=None):
+        h, c = self.init_hidden_state(features)
+        alphas = []
+        captions = []
+        word = torch.tensor(vocab.stoi["<SOS>"]).view(1, -1).to(device)
+        embed = self.embedding(word)
+        for i in range(max_len):
+            alpha, context = self.attention(features, h)
+            alphas.append(alpha.cpu().detach().numpy())
+
+            lstm_inp = torch.cat((embed[:, 0], context), dim=1)
+            h, c = self.lstm(lstm_inp, (h, c))
+            out = self.drop(self.fcn1(h))
+            word_out_idx = torch.argmax(out, dim=1)
+            captions.append(word_out_idx.item())
+            if vocab.itos[word_out_idx.item()] == "<EOS>":
+                break
+            embed = self.embedding(word_out_idx.unsqueeze(0))
+
+        return [vocab.itos[word] for word in captions], alphas
+
+    def load_pretrained_embed(self, embeddings):
+        self.embedding.weight = nn.Parameter(embeddings)
+        for p in self.embedding.parameters():
+            p.requires_grad = True
+
+    def init_hidden_state(self, encoder_output):
+        mean_encoder_out = encoder_output.mean(dim=1)
+        h = self.init_h(mean_encoder_out)
+        c = self.init_c(mean_encoder_out)
+        return h, c
+
+
+class EncoderDecoder(nn.Module):
+    def __init__(self, embed_size, vocab_size, attention_dim, encoder_dim, decoder_dim, fc_dims, p=0.3,
+                 embeddings=None):
+        super().__init__()
+        self.EncoderCNN = Encoder()
+        self.DecoderLSTM = Decoder(embed_size, vocab_size, attention_dim, encoder_dim, decoder_dim, fc_dims, p,
+                                   embeddings)
+
+    def forward(self, imgs, caps):
+        features = self.EncoderCNN(imgs)
+        out = self.DecoderLSTM(features, caps)
+        return out

packages.txt

@@ -0,0 +1 @@
+wget

train.py

@@ -0,0 +1,160 @@
+import torchvision.transforms as T
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data import random_split
+from tqdm import tqdm
+
+from dataset import *
+from model import *
+from utils import *
+
+spacy_eng = spacy.load('en')
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# init seed
+seed = torch.randint(100, (1,))
+torch.manual_seed(seed)
+shuffle = True
+# src folders
+root_folder = "/content/flickr8k/Images"  # change this
+csv_file = "/content/flickr8k/captions.txt"  # change this
+
+# image transforms and augmentation
+transforms = T.Compose([
+    T.Resize(226),
+    T.RandomCrop(224),
+    T.ToTensor(),
+    T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+])
+
+# define dataset
+dataset = FlickrDataset(root_folder, csv_file, transforms)
+
+# split dataset
+val_size = 512
+test_size = 256
+train_size = len(dataset) - val_size - test_size
+train_ds, val_ds, test_ds = random_split(dataset,
+                                         [train_size, val_size, test_size])
+
+# Define data loader parameters
+num_workers = 4
+pin_memory = True
+batch_size_train = 256
+batch_size_val_test = 128
+pad_idx = dataset.vocab.stoi["<PAD>"]
+
+# define loaders
+dataloader_train = DataLoader(train_ds,
+                              batch_size=batch_size_train,
+                              pin_memory=pin_memory,
+                              num_workers=num_workers,
+                              shuffle=shuffle,
+                              collate_fn=CapsCollate(pad_idx=pad_idx, batch_first=True))
+dataloader_validation = DataLoader(val_ds,
+                                   batch_size=batch_size_val_test,
+                                   pin_memory=pin_memory,
+                                   num_workers=num_workers,
+                                   shuffle=shuffle,
+                                   collate_fn=CapsCollate(pad_idx=pad_idx, batch_first=True))
+dataloader_test = DataLoader(test_ds,
+                             batch_size=batch_size_val_test,
+                             pin_memory=pin_memory,
+                             num_workers=num_workers,
+                             shuffle=shuffle,
+                             collate_fn=CapsCollate(pad_idx=pad_idx, batch_first=True))
+
+# model parameters
+embed_wts, embed_size = load_embeding("/content/glove.42B.300d.txt", dataset.vocab)  # change path
+vocab_size = len(dataset.vocab)
+attention_dim = 256
+encoder_dim = 2048
+decoder_dim = 512
+fc_dims = 256
+learning_rate = 5e-4
+
+model = EncoderDecoder(embed_size,
+                       vocab_size,
+                       attention_dim,
+                       encoder_dim,
+                       decoder_dim,
+                       fc_dims,
+                       p=0.3,
+                       embeddings=embed_wts).to(device)
+loss_fn = nn.CrossEntropyLoss(ignore_index=dataset.vocab.stoi["<PAD>"])
+optimizer = optim.Adam(params=model.parameters(), lr=learning_rate)
+
+# training parmeters
+num_epochs = 35
+train_loss_arr = []
+val_loss_arr = []
+
+
+def training(dataset, dataloader, loss_criteria, optimize, grad_clip=5.):
+    total_loss = 0
+    for i, (img, cap) in enumerate(tqdm(dataloader, total=len(dataloader))):
+        img, cap = img.to(device), cap.to(device)
+        optimize.zero_grad()
+        output, attention = model(img, cap)
+        targets = cap[:, 1:]
+        loss = loss_criteria(output.view(-1, vocab_size), targets.reshape(-1))
+        total_loss += (loss.item())
+        loss.backward()
+
+        if grad_clip:
+            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+
+        optimize.step()
+
+    total_loss = total_loss / len(dataloader)
+
+    return total_loss
+
+
+@torch.no_grad()
+def validate(dataset, dataloader, loss_cr):
+    total_loss = 0
+    for val_img, val_cap in tqdm(dataloader, total=len(dataloader)):
+        val_img, val_cap = val_img.to(device), val_cap.to(device)
+        output, attention = model(val_img, val_cap)
+        targets = val_cap[:, 1:]
+        loss = loss_cr(output.view(-1, vocab_size), targets.reshape(-1))
+        total_loss += (loss.item())
+
+    total_loss /= len(dataloader)
+    return total_loss
+
+
+# for see results while training
+@torch.no_grad()
+def test_on_img(data, dataloader):
+    dataiter = iter(dataloader)
+    img, cap = next(dataiter)
+    features = model.EncoderCNN(img[0:1].to(device))
+    caps, alphas = model.DecoderLSTM.gen_captions(features, vocab=data.vocab)
+    caption = ' '.join(caps)
+    show_img(img[0], caption)
+
+
+def main():
+    best_val_loss = 6.0
+    for epoch in range(num_epochs):
+        print(f"Epoch: {epoch + 1}/{num_epochs}")
+        model.train()
+        train_loss = training(dataset, dataloader_train, loss_fn, optimizer)
+        train_loss_arr.append(train_loss)
+
+        model.eval()
+        val_loss = validate(dataset, dataloader_validation, loss_fn)
+        val_loss_arr.append(val_loss)
+        print(f"train_loss: {train_loss} validation_loss: {val_loss}")
+        test_on_img(dataset, dataloader_validation)
+        if len(val_loss_arr) == 1 or val_loss < best_val_loss:
+            best_val_loss = val_loss
+            save_model(model, epoch, optimizer, train_loss, val_loss, vocab=dataset.vocab)
+            print("best model saved successfully")
+
+
+if __name__ == "__main__":
+    print(torch.cuda.is_available())
+    main()

utils.py

@@ -0,0 +1,51 @@
+from matplotlib import pyplot as plt
+import torch
+import numpy as np
+
+
+def show_img(img, caption):
+    img[0] = img[0] * 0.229
+    img[1] = img[1] * 0.224
+    img[2] = img[2] * 0.225
+    img[0] += 0.485
+    img[1] += 0.456
+    img[2] += 0.406
+    img = img.permute(1, 2, 0)
+    img = img.to('cpu').numpy()
+    plt.imshow(img)
+    plt.title(caption)
+    plt.show()
+
+
+def load_embeding(embed_file, vocab):
+    with open(embed_file, 'r') as f:
+        embed_dims = len(f.readline().split(' ')) - 1
+
+    words = set(vocab.stoi.keys())
+    embeddings = torch.FloatTensor(len(words), embed_dims)
+    bias = np.sqrt(3.0 / embeddings.size(1))
+    torch.nn.init.uniform_(embeddings, -bias, bias)
+    print("\nLoading embeddings...")
+    for line in open(embed_file, 'r'):
+        line = line.split(' ')
+        emb_word = line[0]
+        embedding = list(map(lambda t: float(t), filter(lambda n: n and not n.isspace(), line[1:])))
+        # Ignore word if not in train_vocab
+        if emb_word not in words:
+            continue
+        embeddings[vocab.stoi[emb_word]] = torch.FloatTensor(embedding)
+    print("\nEmbeddings loaded!")
+    return embeddings, embed_dims
+
+
+def save_model(model, num_epochs, optimizer, train_loss, val_loss, vocab):
+    model_state = {
+        'num_epochs': num_epochs,
+        'vocab': vocab,
+        'vocab_size': len(vocab.stoi),
+        'state_dict': model.state_dict(),
+        'optimizer_denoise_state_dict': optimizer,
+        'training_loss': train_loss,
+        'val_loss': val_loss,
+    }
+    torch.save(model_state, 'attention_model_state.pth')