def _worker_loop(dataset, index_queue, data_queue, done_event, collate_fn, seed, init_fn, worker_id):
# See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the
# logic of this function.
try:
global _use_shared_memory
_use_shared_memory = True
# Intialize C side signal handlers for SIGBUS and SIGSEGV. Python signal
# module's handlers are executed after Python returns from C low-level
# handlers, likely when the same fatal signal happened again already.
# https://docs.python.org/3/library/signal.html Sec. 18.8.1.1
_set_worker_signal_handlers()
torch.set_num_threads(1)
random.seed(seed)
torch.manual_seed(seed)
data_queue.cancel_join_thread()
if init_fn is not None:
init_fn(worker_id)
watchdog = ManagerWatchdog()
while watchdog.is_alive():
try:
r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
except queue.Empty:
continue
if r is None:
# Received the final signal
assert done_event.is_set()
return
elif done_event.is_set():
# Done event is set. But I haven't received the final signal
# (None) yet. I will keep continuing until get it, and skip the
# processing steps.
continue
idx, batch_indices = r
MyRandomResizedCrop.sample_image_size(idx)
try:
samples = collate_fn([dataset[i] for i in batch_indices])
except Exception:
# It is important that we don't store exc_info in a variable,
# see NOTE [ Python Traceback Reference Cycle Problem ]
data_queue.put((idx, ExceptionWrapper(sys.exc_info())))
else:
data_queue.put((idx, samples))
del samples
except KeyboardInterrupt:
# Main process will raise KeyboardInterrupt anyways.
pass
def build_train_transform(self, image_size=None, print_log=True):
if image_size is None:
image_size = self.image_size
if print_log:
print('Color jitter: %s, resize_scale: %s, img_size: %s' %
(self.distort_color, self.resize_scale, image_size))
if self.distort_color == 'torch':
color_transform = transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1)
elif self.distort_color == 'tf':
color_transform = transforms.ColorJitter(brightness=32. / 255., saturation=0.5)
else:
color_transform = None
if isinstance(image_size, list):
resize_transform_class = MyRandomResizedCrop
print('Use MyRandomResizedCrop: %s, \t %s' % MyRandomResizedCrop.get_candidate_image_size(),
'sync=%s, continuous=%s' % (MyRandomResizedCrop.SYNC_DISTRIBUTED, MyRandomResizedCrop.CONTINUOUS))
else:
resize_transform_class = transforms.RandomResizedCrop
train_transforms = [
resize_transform_class(image_size, scale=(self.resize_scale, 1.0)),
transforms.RandomHorizontalFlip(),
]
if color_transform is not None:
train_transforms.append(color_transform)
train_transforms += [
transforms.ToTensor(),
self.normalize,
]
train_transforms = transforms.Compose(train_transforms)
return train_transforms
def build_train_transform(self,
image_size=None,
print_log=True,
auto_augment='rand-m9-mstd0.5'):
if image_size is None:
image_size = self.image_size
# if print_log:
# print('Color jitter: %s, resize_scale: %s, img_size: %s' %
# (self.distort_color, self.resize_scale, image_size))
# if self.distort_color == 'torch':
# color_transform = transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1)
# elif self.distort_color == 'tf':
# color_transform = transforms.ColorJitter(brightness=32. / 255., saturation=0.5)
# else:
# color_transform = None
if isinstance(image_size, list):
resize_transform_class = MyRandomResizedCrop
print(
'Use MyRandomResizedCrop: %s, \t %s' %
MyRandomResizedCrop.get_candidate_image_size(),
'sync=%s, continuous=%s' %
(MyRandomResizedCrop.SYNC_DISTRIBUTED,
MyRandomResizedCrop.CONTINUOUS))
img_size_min = min(image_size)
else:
resize_transform_class = transforms.RandomResizedCrop
img_size_min = image_size
train_transforms = [
resize_transform_class(image_size, scale=(self.resize_scale, 1.0)),
transforms.RandomHorizontalFlip(),
]
aa_params = dict(
translate_const=int(img_size_min * 0.45),
img_mean=tuple([
min(255, round(255 * x)) for x in
[0.48933587508932375, 0.5183537408957618, 0.5387914411673883]
]),
)
aa_params['interpolation'] = _pil_interp('bicubic')
train_transforms += [rand_augment_transform(auto_augment, aa_params)]
# if color_transform is not None:
# train_transforms.append(color_transform)
train_transforms += [
transforms.ToTensor(),
self.normalize,
]
train_transforms = transforms.Compose(train_transforms)
return train_transforms