/
data.py
201 lines (157 loc) · 7.42 KB
/
data.py
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
from math import ceil
import os
import logging
from pathlib import Path
import json
from PIL import Image
import base64
from io import BytesIO
from dataclasses import dataclass
import lmdb
import pickle
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torchvision.transforms import Compose, Resize, ToTensor, Normalize, InterpolationMode
from timm.data import create_transform
from cn_clip.clip import _tokenizer
from cn_clip.clip import tokenize
def _convert_to_rgb(image):
return image.convert('RGB')
def _preprocess_text(text):
# adapt the text to Chinese BERT vocab
text = text.lower().replace("“", "\"").replace("”", "\"")
return text
class LMDBDataset(Dataset):
def __init__(self, lmdb_path, split="val", max_txt_length=64, use_augment=False, resolution=224):
self.lmdb_path = lmdb_path
# assert LMDB directories exist
assert os.path.isdir(lmdb_path), "The LMDB directory {} of {} split does not exist!".format(lmdb_path, split)
lmdb_pairs = os.path.join(lmdb_path, "pairs")
assert os.path.isdir(lmdb_pairs), "The LMDB directory {} of {} image-text pairs does not exist!".format(lmdb_pairs, split)
lmdb_imgs = os.path.join(lmdb_path, "imgs")
assert os.path.isdir(lmdb_imgs), "The LMDB directory {} of {} image base64 strings does not exist!".format(lmdb_imgs, split)
# open LMDB files
self.env_pairs = lmdb.open(lmdb_pairs, readonly=True, create=False, lock=False, readahead=False, meminit=False)
self.txn_pairs = self.env_pairs.begin(buffers=True)
self.env_imgs = lmdb.open(lmdb_imgs, readonly=True, create=False, lock=False, readahead=False, meminit=False)
self.txn_imgs = self.env_imgs.begin(buffers=True)
# fetch number of pairs and images
self.number_samples = int(self.txn_pairs.get(key=b'num_samples').tobytes().decode('utf-8'))
self.number_images = int(self.txn_imgs.get(key=b'num_images').tobytes().decode('utf-8'))
logging.info("{} LMDB file contains {} images and {} pairs.".format(split, self.number_images, self.number_samples))
super(LMDBDataset, self).__init__()
# the self.dataset_len will be edited to a larger value by calling pad_dataset()
self.dataset_len = self.number_samples
self.global_batch_size = 1 # will be modified to the exact global_batch_size after calling pad_dataset()
self.split = split
self.max_txt_length = max_txt_length
self.use_augment = use_augment
self.transform = self._build_transform(resolution)
def _build_transform(self, resolution):
if self.split == "train" and self.use_augment:
transform = create_transform(
input_size=resolution,
scale=(0.9, 1.0),
is_training=True,
color_jitter=None,
auto_augment='original',
interpolation='bicubic',
mean=(0.48145466, 0.4578275, 0.40821073),
std=(0.26862954, 0.26130258, 0.27577711),
)
transform = Compose(transform.transforms[:-3] + [_convert_to_rgb] + transform.transforms[-3:])
else:
transform = Compose([
Resize((resolution, resolution), interpolation=InterpolationMode.BICUBIC),
_convert_to_rgb,
ToTensor(),
Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
])
return transform
def __del__(self):
if hasattr(self, 'env_pairs'):
self.env_pairs.close()
if hasattr(self, 'env_imgs'):
self.env_imgs.close()
def __len__(self):
return self.dataset_len
def __getitem__(self, index):
sample_index = index % self.number_samples
pair = pickle.loads(self.txn_pairs.get("{}".format(sample_index).encode('utf-8')).tobytes())
image_id, text_id, raw_text = pair
image_b64 = self.txn_imgs.get("{}".format(image_id).encode('utf-8')).tobytes()
image_b64 = image_b64.decode(encoding="utf8", errors="ignore")
image = Image.open(BytesIO(base64.urlsafe_b64decode(image_b64))) # already resized
image = self.transform(image)
text = tokenize([_preprocess_text(raw_text)], context_length=self.max_txt_length)[0]
eos_index = text.numpy().tolist().index(_tokenizer.vocab['[SEP]'])
return image, text, eos_index
def pad_dataset(dataset, global_batch_size):
# edit dataset.__len__() of the dataset
dataset.dataset_len = ceil(dataset.dataset_len / global_batch_size) * global_batch_size
dataset.global_batch_size = global_batch_size
def fetch_resolution(vision_model):
# fetch the resolution from the vision model config
vision_model_config_file = Path(__file__).parent.parent / f"clip/model_configs/{vision_model.replace('/', '-')}.json"
with open(vision_model_config_file, 'r') as fv:
model_info = json.load(fv)
return model_info["image_resolution"]
@dataclass
class DataInfo:
dataloader: DataLoader
sampler: DistributedSampler
dataset: LMDBDataset
epoch_id: int
def get_dataset(args, is_train, max_txt_length=64, epoch_id=0):
if is_train:
db_path = args.train_data
else:
db_path = args.val_data
assert db_path is not None
dataset = LMDBDataset(
db_path,
split="train" if is_train else "val",
max_txt_length=max_txt_length,
use_augment=args.use_augment if is_train else False,
resolution=fetch_resolution(args.vision_model),
)
# pad the dataset splits using the beginning samples in the LMDB files
# to make the number of samples enough for a full final global batch
batch_size = args.batch_size if is_train else args.valid_batch_size
global_batch_size = batch_size * torch.distributed.get_world_size()
pad_dataset(dataset, global_batch_size)
num_samples = dataset.dataset_len
# Update in 22.12.11: We have changed the **validation** dataset sampler during finetuning
# from sequential to shuffled (in a determistic order between experiments and epochs).
# This is to avoid there being one text matching multiple images (or vice versa) in a local batch
# which will affect the correctness of computing the validation in-batch accuracy.
sampler = DistributedSampler(dataset, shuffle=True, seed=args.seed)
sampler.set_epoch(epoch_id if is_train else 0)
dataloader = DataLoader(
dataset,
batch_size=batch_size,
pin_memory=False,
num_workers=args.num_workers if is_train else 1,
sampler=sampler,
)
dataloader.num_samples = num_samples
assert num_samples % dataset.global_batch_size == 0
dataloader.num_batches = num_samples // dataset.global_batch_size
return DataInfo(dataloader, sampler, dataset, epoch_id)
def get_data(args, epoch_id=0, max_txt_length=64):
data = {}
if args.train_data:
data["train"] = get_dataset(
args,
is_train=True,
max_txt_length=max_txt_length,
epoch_id=epoch_id)
if args.val_data:
data["val"] = get_dataset(
args,
is_train=False,
max_txt_length=max_txt_length,
epoch_id=epoch_id)
return data