show-and-tell-resnet-50.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.10","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torchvision.models as models\nimport torchvision.transforms as transforms\nimport numpy as np\n\nfrom os import listdir, makedirs, getcwd, remove\nfrom os.path import isfile, join, abspath, exists, isdir, expanduser\nnp.random.seed(0)\n!ls ../input/pretrained-pytorch-models/","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cache_dir = expanduser('./root/.cache/torch/hub/') #/root/.cache/torch/hub/checkpoints/\nif not exists(cache_dir):\n    makedirs(cache_dir)\nmodels_dir = join(cache_dir, 'checkpoints')\nif not exists(models_dir):\n    makedirs(models_dir)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cp ../input/pretrained-pytorch-models/* /root/.cache/torch/hub/checkpoints/","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls /root/.cache/torch/hub/checkpoints/","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Dataset used - Flickr8k\n# Imports\nimport os  # when loading file paths\nimport pandas as pd  # for lookup in annotation file \"captions.txt\"\nimport spacy  # for tokenizer\nimport torch\nfrom torch.nn.utils.rnn import pad_sequence  # for padding\nfrom torch.utils.data import DataLoader, Dataset\nfrom PIL import Image  # Load img\n\n# Tokenizer\nspacy_eng = spacy.load(\"en\")\n\n\nclass Vocabulary:\n    def __init__(self, freq_threshold):\n        self.itos = {0: \"<PAD>\", 1: \"<SOS>\", 2: \"<EOS>\", 3: \"<UNK>\"}\n        self.stoi = {\"<PAD>\": 0, \"<SOS>\": 1, \"<EOS>\": 2, \"<UNK>\": 3}\n        self.freq_threshold = freq_threshold\n\n    def __len__(self):\n        return len(self.itos)\n\n    @staticmethod\n    def tokenizer_eng(text):\n        return [tok.text.lower() for tok in spacy_eng.tokenizer(text)]\n\n    def build_vocabulary(self, sentence_list):\n        frequencies = {}\n        idx = 4\n        for sentence in sentence_list:\n            for word in self.tokenizer_eng(sentence):\n                if word not in frequencies:\n                    frequencies[word] = 1\n\n                else:\n                    frequencies[word] += 1\n\n                if frequencies[word] == self.freq_threshold:\n                    self.stoi[word] = idx\n                    self.itos[idx] = word\n                    idx += 1\n\n    def numericalize(self, text):\n        tokenized_text = self.tokenizer_eng(text)\n\n        return [\n            self.stoi[token] if token in self.stoi else self.stoi[\"<UNK>\"]\n            for token in tokenized_text\n                ]\n\n\nclass FlickrDataset(Dataset):\n    def __init__(self, root_dir, captions_file, transform=None, freq_threshold=5):\n        self.root_dir = root_dir\n        self.df = pd.read_csv(captions_file)\n        self.transform = transform\n\n        # Get img, caption columns\n        self.imgs = self.df[\"image\"]\n        self.captions = self.df[\"caption\"]\n\n        # Initialize vocabulary and build vocab\n        self.vocab = Vocabulary(freq_threshold)\n        self.vocab.build_vocabulary(self.captions.tolist())\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, index):\n        caption = self.captions[index]\n        img_id = self.imgs[index]\n        img = Image.open(os.path.join(self.root_dir, img_id)).convert(\"RGB\")\n\n        if self.transform is not None:\n            img = self.transform(img)\n\n        numericalized_caption = [self.vocab.stoi[\"<SOS>\"]]\n        numericalized_caption += self.vocab.numericalize(caption)\n        numericalized_caption.append(self.vocab.stoi[\"<EOS>\"])\n\n        return img, torch.tensor(numericalized_caption)\n\n\nclass MyCollate:\n    def __init__(self, pad_idx):\n        self.pad_idx = pad_idx\n\n    def __call__(self, batch):\n        imgs = [item[0].unsqueeze(0) for item in batch]\n        imgs = torch.cat(imgs, dim=0)\n        targets = [item[1] for item in batch]\n        targets = pad_sequence(targets, batch_first=False, padding_value=self.pad_idx)\n\n        return imgs, targets\n\n\ndef get_loader(\n        root_folder,\n        annotation_file,\n        transform,\n        batch_size=32,\n        num_workers=8,\n        shuffle=True,\n        pin_memory=True,):\n    dataset = FlickrDataset(root_folder, annotation_file, transform=transform)\n\n    pad_idx = dataset.vocab.stoi[\"<PAD>\"]\n\n    loader = DataLoader(\n        dataset=dataset,\n        batch_size=batch_size,\n        num_workers=num_workers,\n        shuffle=shuffle,\n        pin_memory=pin_memory,\n        collate_fn=MyCollate(pad_idx=pad_idx),\n    )\n\n    return loader, dataset\nprint(\"Done!\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Imports\nimport torchvision.models as models\nm2 = models.resnet50(pretrained=True)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%whos","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Imports\nimport torchvision.models as models\nresnet = models.resnet50(pretrained=True)\n\n\n# Encoder class of CNN\nclass EncoderCNN(nn.Module):\n    def __init__(self, embed_size, train_CNN=False):\n        super(EncoderCNN, self).__init__()\n        self.train_CNN = train_CNN\n        self.inception = resnet\n        self.inception.fc = nn.Linear(self.inception.fc.in_features, embed_size)\n        self.relu = nn.ReLU()\n        self.times = []\n        self.dropout = nn.Dropout(0.5)\n\n    def forward(self, images):\n        features = self.inception(images)\n        return self.dropout(self.relu(features))\n\n\nclass DecoderRNN(nn.Module):\n    def __init__(self, embed_size, hidden_size, vocab_size, num_layers):\n        super(DecoderRNN, self).__init__()\n        self.embed = nn.Embedding(vocab_size, embed_size)\n        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers)\n        self.linear = nn.Linear(hidden_size, vocab_size)\n        self.dropout = nn.Dropout(0.5)\n\n    def forward(self, features, captions):\n        embeddings = self.dropout(self.embed(captions))\n        embeddings = torch.cat((features.unsqueeze(0), embeddings), dim=0)\n        hiddens, _ = self.lstm(embeddings)\n        outputs = self.linear(hiddens)\n        return outputs\n\n\nclass CNNtoRNN(nn.Module):\n    def __init__(self, embed_size, hidden_size, vocab_size, num_layers):\n        super(CNNtoRNN, self).__init__()\n        self.encoderCNN = EncoderCNN(embed_size)\n        self.decoderRNN = DecoderRNN(embed_size, hidden_size, vocab_size, num_layers)\n\n    def forward(self, images, captions):\n        features = self.encoderCNN(images)\n        outputs = self.decoderRNN(features, captions)\n        return outputs\n\n    def caption_image(self, image, vocabulary, max_length=50):\n        result_caption = []\n\n        with torch.no_grad():\n            x = self.encoderCNN(image).unsqueeze(0)\n            states = None\n\n            for _ in range(max_length):\n                hiddens, states = self.decoderRNN.lstm(x, states)\n                output = self.decoderRNN.linear(hiddens.squeeze(0))\n                predicted = output.argmax(1)\n                result_caption.append(predicted.item())\n                x = self.decoderRNN.embed(predicted).unsqueeze(0)\n\n                if vocabulary.itos[predicted.item()] == \"<EOS>\":\n                    break\n\n        return [vocabulary.itos[idx] for idx in result_caption]\n\ndef predict(model, device, dataset, img):\n    transform = transforms.Compose(\n        [\n            transforms.Resize((299, 299)),\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n        ]\n    )\n\n    model.eval()\n    test_img = transform(Image.open(img).convert(\"RGB\")).unsqueeze(\n        0\n    )\n    return (\"OUTPUT:\"+\" \".join(model.caption_image(test_img.to(device), dataset.vocab)))\n\n\ntransform = transforms.Compose(\n        [\n            transforms.Resize((299, 299)),\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n        ]\n    )\n\nprint(\"Done!\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def save_checkpoint(checkpoint, filename=\"my_checkpoint.pth.tar\"):\n    print(\"==>Saving checkpoint\")\n    torch.save(checkpoint, filename)\n\ndef load_checkpoint(checkpoint, model, optimizer):\n    print(\"==>Loading checkpoint\")\n    model.load_state_dict(checkpoint['state_dict'])\n    optimizer.load_state_dict(checkpoint['optimizer'])\n    ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nfrom tqdm import tqdm\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.transforms as transforms\nfrom torch.utils.tensorboard import SummaryWriter\nfrom PIL import Image\n\ndef train():\n    transform = transforms.Compose(\n        [\n            transforms.Resize((356, 356)),\n            transforms.RandomCrop((299, 299)),\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n        ]\n    )\n\n    train_loader, dataset = get_loader(\n        root_folder=\"../input/flickr8k/Images\",\n        annotation_file=\"../input/flickr8k/captions.txt\",\n        transform=transform,\n        num_workers=2,\n    )\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    train_CNN = False\n\n    # Hyperparameters\n    embed_size = 256\n    hidden_size = 256\n    vocab_size = len(dataset.vocab)\n    num_layers = 1\n    learning_rate = 3e-4\n    num_epochs = 30\n\n    # for tensorboard\n    writer = SummaryWriter(\"runs/flickr\")\n    step = 0\n\n    # initialize model, loss etc\n    model = CNNtoRNN(embed_size, hidden_size, vocab_size, num_layers).to(device)\n    criterion = nn.CrossEntropyLoss(ignore_index=dataset.vocab.stoi[\"<PAD>\"])\n    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n\n    # Only finetune the CNN\n    for name, param in model.encoderCNN.inception.named_parameters():\n        if \"fc.weight\" in name or \"fc.bias\" in name:\n            param.requires_grad = True\n        else:\n            param.requires_grad = train_CNN\n    model.train()\n\n    for epoch in range(num_epochs):\n        if epoch%5 == 0:\n            checkpoint = {'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}\n            save_checkpoint(checkpoint)\n            \n        print(\"Epoch {}\".format(epoch))\n        for idx, (imgs, captions) in tqdm(\n            enumerate(train_loader), total=len(train_loader), leave=False\n        ):\n            imgs = imgs.to(device)\n            captions = captions.to(device)\n\n            outputs = model(imgs, captions[:-1])\n            loss = criterion(\n                outputs.reshape(-1, outputs.shape[2]), captions.reshape(-1)\n            )\n            writer.add_scalar(\"Training loss\", loss.item(), global_step=step)\n            step += 1\n\n            optimizer.zero_grad()\n            loss.backward(loss)\n            optimizer.step()\n            \n    reference = [\n    'A child in a pink dress is climbing up a set of stairs in an entry way .'.split(),\n    'A girl going into a wooden building .'.split(),\n    'A little girl climbing into a wooden playhouse .'.split(),\n    'A little girl climbing the stairs to her playhouse .'.split(),\n    'A little girl in a pink dress going into a wooden cabin'.split()\n                ]\n    candidate=predict(model, device, dataset, img=\"../input/flickr8k/Images/1000268201_693b08cb0e.jpg\")\n    print(candidate)\n    print('BLEU score -> {}'.format(sentence_bleu(reference, candidate.split())))\n    \n\n\nprint(\"Done\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train()\nprint(\"trained successfully\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}