File size: 11,791 Bytes
4d3c316 | 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 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 | import os
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image
from dataset.re_bbox_dataset import re_dataset_bbox, re_eval_dataset
from dataset.randaugment import RandomAugment
#transform 去除了对于图像的水平旋转和镜像
'''
RandomResizedCrop:此操作首先随机调整图像的大小和纵横比,然后对其进行裁剪。结果图像的大小会是config['image_res']。scale=(0.2, 1.0)表示图像的原始面积与裁剪面积之间的比例范围,这里是0.2到1.0之间。interpolation=Image.BICUBIC表示在缩放图像时使用双三次插值。
RandomHorizontalFlip:以50%的概率对图像进行水平翻转。
RandomAugment:这是一种数据增强方法,但并不是torchvision库中的标准方法,可能是某个特定库或自定义实现的。从其参数可以看出,它将在给定的数据增强方法列表augs中随机选择2种方法,并应用最多7次的随机增强。所列出的增强方法包括:
Identity:保持不变
AutoContrast:自动调整图像的对比度
Equalize:使图像的直方图均衡化
Brightness:调整图像的亮度
Sharpness:调整图像的锐度
ShearX、ShearY:对图像进行X轴或Y轴的扭曲
TranslateX、TranslateY:对图像在X轴或Y轴上进行随机平移
Rotate 随机旋转图像
ToTensor 将PIL Image或numpy ndarray的图像数据转换为torch.Tensor 并且调整其范围从[0, 255]到[0.0, 1.0]。
normalize 对图像数据进行归一化。具体的均值和标准差应该在normalize变量中定义 但您没有提供这部分代码。归一化有助于神经网络的训练稳定和加速。
'''
def create_dataset(dataset, config, evaluate=False):
normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
pretrain_transform = transforms.Compose([
transforms.RandomResizedCrop(config['image_res'], scale=(0.2, 1.0),
interpolation=Image.BICUBIC),
# transforms.RandomHorizontalFlip(),
RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
transforms.ToTensor(),
normalize,
])
train_transform = transforms.Compose([
transforms.RandomResizedCrop(config['image_res'], scale=(0.5, 1.0),
interpolation=Image.BICUBIC),
# transforms.RandomHorizontalFlip(),
RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
transforms.ToTensor(),
normalize,
])
train_transform_wohflip = transforms.Compose([
transforms.RandomResizedCrop(config['image_res'], scale=(0.5, 1.0),
interpolation=Image.BICUBIC),
# transforms.RandomHorizontalFlip(),
RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness',]),
transforms.ToTensor(),
normalize,
])
box_transform = transforms.Compose([
RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
transforms.ToTensor(),
normalize,
])
test_transform = transforms.Compose([
transforms.Resize((config['image_res'], config['image_res']), interpolation=Image.BICUBIC),
transforms.ToTensor(),
normalize,
])
if dataset == 'pretrain':
general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
transform=pretrain_transform)
region_dataset = RegionTextJsonDataset(config, config['train_file_regions'], rank=int(os.environ.get('RANK') or 0),
world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
transform=pretrain_transform, box_transform=box_transform)
return general_dataset, region_dataset
elif dataset == 're':
test_dataset = re_eval_dataset(config['test_file'], test_transform, config['image_root'])
if evaluate:
return None, test_dataset
train_dataset = re_train_dataset(config['train_file'], train_transform, config['image_root'])
# val_dataset = re_eval_dataset(config['val_file'], test_transform, config['image_root'])
return train_dataset, test_dataset
#修改
elif dataset == 're_bbox':
test_dataset = re_eval_dataset(config['test_file'], test_transform, config['image_root'])
if evaluate:
return None, test_dataset
train_transform = transforms.Compose([
RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
transforms.ToTensor(),
normalize,
])
train_dataset = re_dataset_bbox(config['train_file'], train_transform, config['image_root'], mode='train', config=config)
return train_dataset, test_dataset
elif dataset == 'vqa':
vqa_test_dataset = vqa_dataset(config['test_file'], test_transform, config['vqa_root'], config['vg_root'],
split='test', answer_list=config['answer_list'],
text_encoder=config['text_encoder'], use_roberta=config['use_roberta'])
if evaluate:
return None, vqa_test_dataset
train_dataset = vqa_dataset(config['train_file'], train_transform_wohflip, config['vqa_root'], config['vg_root'],
split='train', text_encoder=config['text_encoder'], use_roberta=config['use_roberta'])
return train_dataset, vqa_test_dataset
elif dataset == 'nlvr_pretrain':
general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
transform=pretrain_transform)
return general_dataset
elif dataset == 'nlvr':
test_dataset = nlvr_dataset(config['test_file'], test_transform, config['image_root'])
if evaluate:
return None, None, test_dataset
train_dataset = nlvr_dataset(config['train_file'], train_transform, config['image_root'])
val_dataset = nlvr_dataset(config['val_file'], test_transform, config['image_root'])
return train_dataset, val_dataset, test_dataset
elif dataset == 'grounding':
test_dataset = grounding_dataset(config['test_file'], test_transform, config['image_root'], mode='test')
if evaluate:
return None, test_dataset
train_transform = transforms.Compose([
transforms.Resize((config['image_res'], config['image_res']), interpolation=Image.BICUBIC),
transforms.RandomHorizontalFlip(),
RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
transforms.ToTensor(),
normalize,
])
train_dataset = grounding_dataset(config['train_file'], train_transform, config['image_root'], mode='train')
return train_dataset, test_dataset
elif dataset == 'grounding_bbox_pretrain':
region_dataset = RegionTextJsonDataset(config, config['train_file_regions'], rank=int(os.environ.get('RANK') or 0),
world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
transform=pretrain_transform, box_transform=box_transform)
return region_dataset
elif dataset == 'grounding_bbox':
test_dataset = grounding_dataset_bbox(config['test_file'], test_transform, config['image_root'], mode='test', config=config)
if evaluate:
return None, test_dataset
train_transform = transforms.Compose([
RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
transforms.ToTensor(),
normalize,
])
train_dataset = grounding_dataset_bbox(config['train_file'], train_transform, config['image_root'], mode='train', config=config)
return train_dataset, test_dataset
elif dataset == 'captioning_pretrain':
general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
transform=pretrain_transform, add_eos=True)
return general_dataset
elif dataset == 'caption_coco':
train_dataset = coco_karpathy_train(train_transform, config['image_root'], config['train_file'], prompt=config['prompt'], max_words=config['max_tokens'])
val_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['val_file'], 'val')
test_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['test_file'], 'test')
return train_dataset, val_dataset, test_dataset
elif dataset == 'caption_coco_scst':
train_dataset = coco_karpathy_train_scst(train_transform, config['image_root'], config['train_file'],
prompt=config['prompt'], max_words=config['max_tokens'])
val_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['val_file'], 'val')
test_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['test_file'], 'test')
return train_dataset, val_dataset, test_dataset
else:
raise NotImplementedError(f"dataset == {dataset}")
def vqa_collate_fn(batch):
image_list, question_list, answer_list, weight_list, n = [], [], [], [], []
for image, question, answer, weights in batch:
image_list.append(image)
question_list.append(question)
weight_list += weights
answer_list += answer
n.append(len(answer))
return torch.stack(image_list, dim=0), question_list, answer_list, torch.Tensor(weight_list), n
def create_sampler(datasets, shuffles, num_tasks, global_rank):
samplers = []
for dataset, shuffle in zip(datasets, shuffles):
sampler = torch.utils.data.DistributedSampler(dataset, num_replicas=num_tasks, rank=global_rank, shuffle=shuffle)
samplers.append(sampler)
return samplers
def create_loader(datasets, samplers, batch_size, num_workers, is_trains, collate_fns):
loaders = []
for dataset, sampler, bs, n_worker, is_train, collate_fn in zip(datasets, samplers, batch_size, num_workers,
is_trains, collate_fns):
if is_train:
shuffle = (sampler is None)
drop_last = True
else:
shuffle = False
drop_last = False
loader = DataLoader(
dataset,
batch_size=bs,
num_workers=n_worker,
pin_memory=True,
sampler=sampler,
shuffle=shuffle,
collate_fn=collate_fn,
drop_last=drop_last,
)
loaders.append(loader)
if len(loaders) <= 1:
print(f"### be careful: func create_loader returns a list length of {len(loaders)}")
return loaders
|