def test_autoaugment(self):
tensor = torch.randint(0,
256,
size=(3, 44, 56),
dtype=torch.uint8,
device=self.device)
batch_tensors = torch.randint(0,
256,
size=(4, 3, 44, 56),
dtype=torch.uint8,
device=self.device)
s_transform = None
for policy in T.AutoAugmentPolicy:
for fill in [
None, 85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [
1,
], 1
]:
transform = T.AutoAugment(policy=policy, fill=fill)
s_transform = torch.jit.script(transform)
for _ in range(25):
_test_transform_vs_scripted(transform, s_transform, tensor)
_test_transform_vs_scripted_on_batch(
transform, s_transform, batch_tensors)
if s_transform is not None:
with get_tmp_dir() as tmp_dir:
s_transform.save(os.path.join(tmp_dir, "t_autoaugment.pt"))
def __init__(self,
root: str,
image_size: int = 128,
upscale_factor: int = 4):
r"""
Args:
root (str): The directory address where the data image is stored.
image_size (optional, int): The size of image block is randomly cut out from the original image. (Default: 128)
upscale_factor (optional, int): Image magnification. (Default: 4)
"""
super(BaseTestDataset, self).__init__()
self.filenames = [
os.path.join(root, x) for x in os.listdir(root)
if check_image_file(x)
]
self.lr_transforms = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(
(image_size // upscale_factor, image_size // upscale_factor),
interpolation=InterpolationMode.BICUBIC),
transforms.ToTensor()
])
self.bicubic_transforms = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((image_size, image_size),
interpolation=InterpolationMode.BICUBIC),
transforms.ToTensor()
])
self.hr_transforms = transforms.Compose([
transforms.RandomCrop((image_size, image_size)),
transforms.AutoAugment(),
transforms.ToTensor()
])
def test_autoaugment(device, policy, fill):
tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
transform = T.AutoAugment(policy=policy, fill=fill)
s_transform = torch.jit.script(transform)
for _ in range(25):
_test_transform_vs_scripted(transform, s_transform, tensor)
_test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
def cars_train_transfroms_autoaugment(img_size:int=448,
mean:list=IMAGENET_DEFAULT_MEAN,
std:list=IMAGENET_DEFAULT_STD,
interpolation=Image.BILINEAR,
):
transform=transforms.Compose([
transforms.Resize((600, 600), interpolation),
transforms.RandomCrop((img_size, img_size)),
transforms.RandomHorizontalFlip(),
transforms.AutoAugment(),
transforms.ToTensor(),
transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)]
)
return transform
def main():
# Create a pytorch dataset
data_dir = pathlib.Path('./data/tiny-imagenet-200')
image_count = len(list(data_dir.glob('**/*.JPEG')))
CLASS_NAMES = np.array(
[item.name for item in (data_dir / 'train').glob('*')])
print('Discovered {} images'.format(image_count))
# Create the training data generator
batch_size = 32
im_height = 64
im_width = 64
num_epochs = 1
data_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.AutoAugment(transforms.AutoAugmentPolicy.IMAGENET),
transforms.Resize((Net.input_dim, Net.input_dim)),
transforms.Normalize((0, 0, 0), tuple(np.sqrt((255, 255, 255)))),
])
train_set = torchvision.datasets.ImageFolder(data_dir / 'train',
data_transforms)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=4,
pin_memory=True)
# Create a simple model
model = Net(len(CLASS_NAMES), im_height, im_width)
optim = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
for i in range(num_epochs):
train_total, train_correct = 0, 0
for idx, (inputs, targets) in enumerate(train_loader):
optim.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optim.step()
_, predicted = outputs.max(1)
train_total += targets.size(0)
train_correct += predicted.eq(targets).sum().item()
print("\r", end='')
print(
f'training {100 * idx / len(train_loader):.2f}%: {train_correct / train_total:.3f}',
end='')
torch.save({
'net': model.state_dict(),
}, 'latest.pt')
def get_transform_cifar(
mode: str,
auto_augment: bool = False,
cutout: bool = False,
cutout_length: int = None,
data_shape: list[int] = [3, 32, 32]) -> transforms.Compose:
if mode != 'train':
return transforms.Compose([
transforms.PILToTensor(),
transforms.ConvertImageDtype(torch.float)
])
cutout_length = cutout_length or data_shape[-1] // 2
transform_list = [
transforms.RandomCrop(data_shape[-2:], padding=data_shape[-1] // 8),
transforms.RandomHorizontalFlip(),
]
if auto_augment:
transform_list.append(
transforms.AutoAugment(transforms.AutoAugmentPolicy.CIFAR10))
transform_list.append(transforms.PILToTensor())
transform_list.append(transforms.ConvertImageDtype(torch.float))
if cutout:
transform_list.append(Cutout(cutout_length))
return transforms.Compose(transform_list)
def get_transform_imagenet(mode: str,
use_tuple: bool = False,
auto_augment: bool = False) -> transforms.Compose:
if mode == 'train':
transform_list = [
transforms.RandomResizedCrop((224, 224) if use_tuple else 224),
transforms.RandomHorizontalFlip(),
# transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4), # noqa
]
if auto_augment:
transform_list.append(
transforms.AutoAugment(transforms.AutoAugmentPolicy.IMAGENET))
transform_list.append(transforms.PILToTensor())
transform_list.append(transforms.ConvertImageDtype(torch.float))
transform = transforms.Compose(transform_list)
else:
# TODO: torchvision.prototypes.transforms._presets.ImageClassificationEval
transform = transforms.Compose([
transforms.Resize((256, 256) if use_tuple else 256),
transforms.CenterCrop((224, 224) if use_tuple else 224),
transforms.PILToTensor(),
transforms.ConvertImageDtype(torch.float)
])
return transform
# size.
resize_cropper = T.RandomResizedCrop(size=(32, 32))
resized_crops = [resize_cropper(orig_img) for _ in range(4)]
plot(resized_crops)
####################################
# AutoAugment
# ~~~~~~~~~~~
# The :class:`~torchvision.transforms.AutoAugment` transform
# automatically augments data based on a given auto-augmentation policy.
# See :class:`~torchvision.transforms.AutoAugmentPolicy` for the available policies.
policies = [
T.AutoAugmentPolicy.CIFAR10, T.AutoAugmentPolicy.IMAGENET,
T.AutoAugmentPolicy.SVHN
]
augmenters = [T.AutoAugment(policy) for policy in policies]
imgs = [[augmenter(orig_img) for _ in range(4)] for augmenter in augmenters]
row_title = [str(policy).split('.')[-1] for policy in policies]
plot(imgs, row_title=row_title)
####################################
# Randomly-applied transforms
# ---------------------------
#
# Some transforms are randomly-applied given a probability ``p``. That is, the
# transformed image may actually be the same as the original one, even when
# called with the same transformer instance!
#
# RandomHorizontalFlip
# ~~~~~~~~~~~~~~~~~~~~
# The :class:`~torchvision.transforms.RandomHorizontalFlip` transform
def test_autoaugment_save():
transform = T.AutoAugment()
s_transform = torch.jit.script(transform)
with get_tmp_dir() as tmp_dir:
s_transform.save(os.path.join(tmp_dir, "t_autoaugment.pt"))
def load_data(traindir, valdir, args):
# Data loading code
print("Loading data")
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
print("Loading training data")
st = time.time()
cache_path = _get_cache_path(traindir)
if args.cache_dataset and os.path.exists(cache_path):
# Attention, as the transforms are also cached!
print("Loading dataset_train from {}".format(cache_path))
dataset, _ = torch.load(cache_path)
else:
trans = [
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
]
if args.auto_augment is not None:
aa_policy = transforms.AutoAugmentPolicy(args.auto_augment)
trans.append(transforms.AutoAugment(policy=aa_policy))
trans.extend([
transforms.ToTensor(),
normalize,
])
if args.random_erase > 0:
trans.append(transforms.RandomErasing(p=args.random_erase))
dataset = torchvision.datasets.ImageFolder(
traindir,
transforms.Compose(trans))
if args.cache_dataset:
print("Saving dataset_train to {}".format(cache_path))
utils.mkdir(os.path.dirname(cache_path))
utils.save_on_master((dataset, traindir), cache_path)
print("Took", time.time() - st)
print("Loading validation data")
cache_path = _get_cache_path(valdir)
if args.cache_dataset and os.path.exists(cache_path):
# Attention, as the transforms are also cached!
print("Loading dataset_test from {}".format(cache_path))
dataset_test, _ = torch.load(cache_path)
else:
dataset_test = torchvision.datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
if args.cache_dataset:
print("Saving dataset_test to {}".format(cache_path))
utils.mkdir(os.path.dirname(cache_path))
utils.save_on_master((dataset_test, valdir), cache_path)
print("Creating data loaders")
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test)
else:
train_sampler = torch.utils.data.RandomSampler(dataset)
test_sampler = torch.utils.data.SequentialSampler(dataset_test)
return dataset, dataset_test, train_sampler, test_sampler
def main():
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))
# Data transformations
trans_train = transforms.Compose([
transforms.RandomApply(transforms=[
transforms.AutoAugment(transforms.AutoAugmentPolicy.IMAGENET),
# transforms.RandomPerspective(distortion_scale=0.6, p=1.0),
transforms.RandomRotation(degrees=(0, 180)),
transforms.RandomHorizontalFlip(),
]),
np.float32,
transforms.ToTensor(),
fixed_image_standardization,
])
trans_val = transforms.Compose([
# transforms.CenterCrop(120),
np.float32,
transforms.ToTensor(),
fixed_image_standardization,
])
train_dataset = datasets.ImageFolder(os.path.join(data_dir,
"train_aligned"),
transform=trans_train)
val_dataset = datasets.ImageFolder(os.path.join(data_dir, "val_aligned"),
transform=trans_val)
# Prepare the model
model = InceptionResnetV1(classify=False,
pretrained="vggface2",
dropout_prob=0.5).to(device)
# for param in list(model.parameters())[:-8]:
# param.requires_grad = False
trunk_optimizer = torch.optim.SGD(model.parameters(), lr=LR)
# Set the loss function
loss = losses.ArcFaceLoss(len(train_dataset.classes), 512)
# Package the above stuff into dictionaries.
models = {"trunk": model}
optimizers = {"trunk_optimizer": trunk_optimizer}
loss_funcs = {"metric_loss": loss}
mining_funcs = {}
lr_scheduler = {
"trunk_scheduler_by_plateau":
torch.optim.lr_scheduler.ReduceLROnPlateau(trunk_optimizer)
}
# Create the tester
record_keeper, _, _ = logging_presets.get_record_keeper(
"logs", "tensorboard")
hooks = logging_presets.get_hook_container(record_keeper)
dataset_dict = {"val": val_dataset, "train": train_dataset}
model_folder = "training_saved_models"
def visualizer_hook(umapper, umap_embeddings, labels, split_name, keyname,
*args):
logging.info("UMAP plot for the {} split and label set {}".format(
split_name, keyname))
label_set = np.unique(labels)
num_classes = len(label_set)
fig = plt.figure(figsize=(8, 7))
plt.gca().set_prop_cycle(
cycler("color", [
plt.cm.nipy_spectral(i)
for i in np.linspace(0, 0.9, num_classes)
]))
for i in range(num_classes):
idx = labels == label_set[i]
plt.plot(umap_embeddings[idx, 0],
umap_embeddings[idx, 1],
".",
markersize=1)
plt.show()
tester = testers.GlobalEmbeddingSpaceTester(
end_of_testing_hook=hooks.end_of_testing_hook,
dataloader_num_workers=4,
accuracy_calculator=AccuracyCalculator(
include=['mean_average_precision_at_r'], k="max_bin_count"))
end_of_epoch_hook = hooks.end_of_epoch_hook(tester,
dataset_dict,
model_folder,
splits_to_eval=[('val',
['train'])])
# Create the trainer
trainer = trainers.MetricLossOnly(
models,
optimizers,
batch_size,
loss_funcs,
mining_funcs,
train_dataset,
lr_schedulers=lr_scheduler,
dataloader_num_workers=8,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=end_of_epoch_hook)
trainer.train(num_epochs=num_epochs)
# The :class:`~torchvision.transforms.GaussianBlur` transform
# (see also :func:`~torchvision.transforms.functional.gaussian_blur`)
# performs gaussianblur transform on an image.
gaus_blur_img = T.GaussianBlur(kernel_size=(5, 9), sigma=(0.4, 3.0))(orig_img)
plot(gaus_blur_img, "Gaussian Blurred Image")
####################################
# AutoAugment
# -----------
# The :class:`~torchvision.transforms.AutoAugment` transform
# automatically augments data based on a given auto-augmentation policy.
# See :class:`~torchvision.transforms.AutoAugmentPolicy` for the available policies.
policies = [
T.AutoAugmentPolicy.CIFAR10, T.AutoAugmentPolicy.IMAGENET,
T.AutoAugmentPolicy.SVHN
]
num_cols = 5
fig, axs = plt.subplots(nrows=len(policies), ncols=num_cols)
fig.suptitle("Auto-augmented images with different policies")
for pol_idx, policy in enumerate(policies):
auto_augmenter = T.AutoAugment(policy)
for col in range(num_cols):
augmented_img = auto_augmenter(orig_img)
ax = axs[pol_idx, col]
ax.imshow(np.asarray(augmented_img))
ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
if col == 0:
ax.set(ylabel=str(policy).split('.')[-1])
