Views
No views yet
Matthijs/snacks datasettorchvision.set_transform method.1from transformers import ViTFeatureExtractor
2from torchvision.transforms import (
3 Compose,
4 Normalize,
5 Resize,
6 RandomResizedCrop,
7 RandomHorizontalFlip,
8 RandomAdjustSharpness,
9 ToTensor
10)
11
12checkpoint = 'google/vit-base-patch16-224-in21k'
13feature_extractor = ViTFeatureExtractor.from_pretrained(checkpoint)
14
15# transformations on the training set
16train_aug_transforms = Compose([
17 RandomResizedCrop(size=feature_extractor.size),
18 RandomHorizontalFlip(p=0.5),
19 RandomAdjustSharpness(sharpness_factor=5, p=0.5),
20 ToTensor(),
21 Normalize(mean=feature_extractor.image_mean, std=feature_extractor.image_std),
22])
23
24# transformations on the validation/test set
25valid_aug_transforms = Compose([
26 Resize(size=(feature_extractor.size, feature_extractor.size)),
27 ToTensor(),
28 Normalize(mean=feature_extractor.image_mean, std=feature_extractor.image_std),
29])