Spaces:
Runtime error
Runtime error
File size: 3,918 Bytes
aa2ec08 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import pandas as pd
import numpy as np
from collections import Counter
import torchvision
from torchvision import transforms
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torch
import torch.nn as nn
from torch.autograd import Variable
from tqdm.notebook import tqdm
import matplotlib.pyplot as plt
import math
import torch.nn.functional as F
import pickle
import gc
import random
pd.set_option('display.max_colwidth', None)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
max_seq_len = 33
vocab_size = 8360
class PositionalEncoding(nn.Module):
def __init__(self, d_model, dropout=0.1, max_len=max_seq_len):
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)
self.register_buffer('pe', pe)
def forward(self, x):
if self.pe.size(0) < x.size(0):
self.pe = self.pe.repeat(x.size(0), 1, 1).to(device)
self.pe = self.pe[:x.size(0), : , : ]
x = x + self.pe
return self.dropout(x)
class ImageCaptionModel(nn.Module):
def __init__(self, n_head, n_decoder_layer, vocab_size, embedding_size):
super(ImageCaptionModel, self).__init__()
self.pos_encoder = PositionalEncoding(embedding_size, 0.1)
self.TransformerDecoderLayer = nn.TransformerDecoderLayer(d_model = embedding_size, nhead = n_head)
self.TransformerDecoder = nn.TransformerDecoder(decoder_layer = self.TransformerDecoderLayer, num_layers = n_decoder_layer)
self.embedding_size = embedding_size
self.embedding = nn.Embedding(vocab_size , embedding_size)
self.last_linear_layer = nn.Linear(embedding_size, vocab_size)
self.init_weights()
def init_weights(self):
initrange = 0.1
self.embedding.weight.data.uniform_(-initrange, initrange)
self.last_linear_layer.bias.data.zero_()
self.last_linear_layer.weight.data.uniform_(-initrange, initrange)
def generate_Mask(self, size, decoder_inp):
decoder_input_mask = (torch.triu(torch.ones(size, size)) == 1).transpose(0, 1)
decoder_input_mask = decoder_input_mask.float().masked_fill(decoder_input_mask == 0, float('-inf')).masked_fill(decoder_input_mask == 1, float(0.0))
decoder_input_pad_mask = decoder_inp.float().masked_fill(decoder_inp == 0, float(0.0)).masked_fill(decoder_inp > 0, float(1.0))
decoder_input_pad_mask_bool = decoder_inp == 0
return decoder_input_mask, decoder_input_pad_mask, decoder_input_pad_mask_bool
def forward(self, encoded_image, decoder_inp):
encoded_image = encoded_image.permute(1,0,2)
decoder_inp_embed = self.embedding(decoder_inp)* math.sqrt(self.embedding_size)
decoder_inp_embed = self.pos_encoder(decoder_inp_embed)
decoder_inp_embed = decoder_inp_embed.permute(1,0,2)
decoder_input_mask, decoder_input_pad_mask, decoder_input_pad_mask_bool = self.generate_Mask(decoder_inp.size(1), decoder_inp)
decoder_input_mask = decoder_input_mask.to(device)
decoder_input_pad_mask = decoder_input_pad_mask.to(device)
decoder_input_pad_mask_bool = decoder_input_pad_mask_bool.to(device)
decoder_output = self.TransformerDecoder(tgt = decoder_inp_embed, memory = encoded_image, tgt_mask = decoder_input_mask, tgt_key_padding_mask = decoder_input_pad_mask_bool)
final_output = self.last_linear_layer(decoder_output)
return final_output, decoder_input_pad_mask
|