chendl/compositional_test

1# coding=utf-82# Copyright (c) Facebook, Inc. and its affiliates.3# Copyright (c) HuggingFace Inc. team.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9#     http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16 17import json18import os19from collections import Counter20 21import torch22import torchvision23import torchvision.transforms as transforms24from PIL import Image25from torch import nn26from torch.utils.data import Dataset27 28 29POOLING_BREAKDOWN = {1: (1, 1), 2: (2, 1), 3: (3, 1), 4: (2, 2), 5: (5, 1), 6: (3, 2), 7: (7, 1), 8: (4, 2), 9: (3, 3)}30 31 32class ImageEncoder(nn.Module):33    def __init__(self, args):34        super().__init__()35        model = torchvision.models.resnet152(pretrained=True)36        modules = list(model.children())[:-2]37        self.model = nn.Sequential(*modules)38        self.pool = nn.AdaptiveAvgPool2d(POOLING_BREAKDOWN[args.num_image_embeds])39 40    def forward(self, x):41        # Bx3x224x224 -> Bx2048x7x7 -> Bx2048xN -> BxNx204842        out = self.pool(self.model(x))43        out = torch.flatten(out, start_dim=2)44        out = out.transpose(1, 2).contiguous()45        return out  # BxNx204846 47 48class JsonlDataset(Dataset):49    def __init__(self, data_path, tokenizer, transforms, labels, max_seq_length):50        self.data = [json.loads(l) for l in open(data_path)]51        self.data_dir = os.path.dirname(data_path)52        self.tokenizer = tokenizer53        self.labels = labels54        self.n_classes = len(labels)55        self.max_seq_length = max_seq_length56 57        self.transforms = transforms58 59    def __len__(self):60        return len(self.data)61 62    def __getitem__(self, index):63        sentence = torch.LongTensor(self.tokenizer.encode(self.data[index]["text"], add_special_tokens=True))64        start_token, sentence, end_token = sentence[0], sentence[1:-1], sentence[-1]65        sentence = sentence[: self.max_seq_length]66 67        label = torch.zeros(self.n_classes)68        label[[self.labels.index(tgt) for tgt in self.data[index]["label"]]] = 169 70        image = Image.open(os.path.join(self.data_dir, self.data[index]["img"])).convert("RGB")71        image = self.transforms(image)72 73        return {74            "image_start_token": start_token,75            "image_end_token": end_token,76            "sentence": sentence,77            "image": image,78            "label": label,79        }80 81    def get_label_frequencies(self):82        label_freqs = Counter()83        for row in self.data:84            label_freqs.update(row["label"])85        return label_freqs86 87 88def collate_fn(batch):89    lens = [len(row["sentence"]) for row in batch]90    bsz, max_seq_len = len(batch), max(lens)91 92    mask_tensor = torch.zeros(bsz, max_seq_len, dtype=torch.long)93    text_tensor = torch.zeros(bsz, max_seq_len, dtype=torch.long)94 95    for i_batch, (input_row, length) in enumerate(zip(batch, lens)):96        text_tensor[i_batch, :length] = input_row["sentence"]97        mask_tensor[i_batch, :length] = 198 99    img_tensor = torch.stack([row["image"] for row in batch])100    tgt_tensor = torch.stack([row["label"] for row in batch])101    img_start_token = torch.stack([row["image_start_token"] for row in batch])102    img_end_token = torch.stack([row["image_end_token"] for row in batch])103 104    return text_tensor, mask_tensor, img_tensor, img_start_token, img_end_token, tgt_tensor105 106 107def get_mmimdb_labels():108    return [109        "Crime",110        "Drama",111        "Thriller",112        "Action",113        "Comedy",114        "Romance",115        "Documentary",116        "Short",117        "Mystery",118        "History",119        "Family",120        "Adventure",121        "Fantasy",122        "Sci-Fi",123        "Western",124        "Horror",125        "Sport",126        "War",127        "Music",128        "Musical",129        "Animation",130        "Biography",131        "Film-Noir",132    ]133 134 135def get_image_transforms():136    return transforms.Compose(137        [138            transforms.Resize(256),139            transforms.CenterCrop(224),140            transforms.ToTensor(),141            transforms.Normalize(142                mean=[0.46777044, 0.44531429, 0.40661017],143                std=[0.12221994, 0.12145835, 0.14380469],144            ),145        ]146    )147