How to use the albumentations.Normalize function in albumentations

# FOR DISTRIBUTED: Initialize the backend. torch.distributed.launch will
        # provide environment variables, and requires that you use init_method=`env://`.
        torch.distributed.init_process_group(backend="nccl", init_method="env://")

    scheduler_step = config.TRAIN.END_EPOCH // config.TRAIN.SNAPSHOTS
    torch.backends.cudnn.benchmark = config.CUDNN.BENCHMARK

    torch.manual_seed(config.SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(config.SEED)
    np.random.seed(seed=config.SEED)
    # Setup Augmentations
    basic_aug = Compose(
        [
            Normalize(mean=(config.TRAIN.MEAN,), std=(config.TRAIN.STD,), max_pixel_value=1),
            Resize(
                config.TRAIN.AUGMENTATIONS.RESIZE.HEIGHT, config.TRAIN.AUGMENTATIONS.RESIZE.WIDTH, always_apply=True,
            ),
            PadIfNeeded(
                min_height=config.TRAIN.AUGMENTATIONS.PAD.HEIGHT,
                min_width=config.TRAIN.AUGMENTATIONS.PAD.WIDTH,
                border_mode=cv2.BORDER_CONSTANT,
                always_apply=True,
                mask_value=255,
            ),
        ]
    )
    if config.TRAIN.AUGMENTATION:
        train_aug = Compose([basic_aug, HorizontalFlip(p=0.5)])
        val_aug = basic_aug
    else:

def train(dataset_dir, batch_size, learning_rate, epochs):

    input_paths, target_labels = prepare_samples(dataset_dir, train_or_test='train')
    preprocess = A.Compose([A.Normalize()])
    augmentation = None
    dataset = Cifar10(input_paths, target_labels, batch_size, shuffle=True,
                      preprocess=preprocess, augmentation=augmentation)

    model = get_resnet_mini(input_shape=(32, 32, 3),
                            num_classes=10)
    model.summary()

    optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    metric = tf.keras.metrics.SparseCategoricalAccuracy()

    @tf.function
    def train_step(x, y):
        with tf.GradientTape() as tape:
            logits = model(x, training=True)

def get_test_transforms():
        return A.Compose([A.Normalize(mean=(0.485, ), std=(0.229, ))])

decoder_merge_policy="cat",
)
pad_factor = 64
imread_library = "cv2"  # can be cv2 or jpeg4py

optimizer = Adam(
    [
        {"params": model.decoder.parameters(), "lr": train_parameters["lr"]},
        # decrease lr for encoder in order not to permute
        # pre-trained weights with large gradients on training start
        {"params": model.encoder.parameters(), "lr": train_parameters["lr"] / 100},
    ],
    weight_decay=1e-4,
)

normalization = albu.Normalize(mean=mean, std=std, p=1)

train_augmentations = albu.Compose(
    [
        albu.RandomSizedCrop(
            min_max_height=(
                int(0.5 * (train_parameters["height_crop_size"])),
                int(2 * (train_parameters["height_crop_size"])),
            ),
            height=train_parameters["height_crop_size"],
            width=train_parameters["width_crop_size"],
            w2h_ratio=1.0,
            p=1,
        ),
        albu.HorizontalFlip(p=0.5),
        normalization,
    ],

)

pad_factor = 64
imread_library = "cv2"  # can be cv2 or jpeg4py

optimizer = RAdam(
    [
        {"params": model.decoder.parameters(), "lr": train_parameters["lr"]},
        # decrease lr for encoder in order not to permute
        # pre-trained weights with large gradients on training start
        {"params": model.encoder.parameters(), "lr": train_parameters["lr"] / 100},
    ],
    weight_decay=1e-3,
)

normalization = albu.Normalize(mean=mean, std=std, p=1)

train_augmentations = albu.Compose(
    [
        albu.RandomSizedCrop(
            min_max_height=(
                int(0.5 * (train_parameters["height_crop_size"])),
                int(2 * (train_parameters["height_crop_size"])),
            ),
            height=train_parameters["height_crop_size"],
            width=train_parameters["width_crop_size"],
            w2h_ratio=1.0,
            p=1,
        ),
        albu.ShiftScaleRotate(
            border_mode=cv2.BORDER_CONSTANT, rotate_limit=10, scale_limit=0, p=0.5, mask_value=ignore_index
        ),

std = [0.229, 0.224, 0.225]

train_transforms = A.Compose([
    A.RandomResizedCrop(train_crop_size, train_crop_size, scale=(0.08, 1.0)),
    A.HorizontalFlip(),
    A.CoarseDropout(max_height=32, max_width=32),
    A.HueSaturationValue(),
    A.Normalize(mean=mean, std=std),
    ToTensor(),
])

val_transforms = A.Compose([
    # https://github.com/facebookresearch/FixRes/blob/b27575208a7c48a3a6e0fa9efb57baa4021d1305/imnet_resnet50_scratch/transforms.py#L76
    A.Resize(int((256 / 224) * val_crop_size), int((256 / 224) * val_crop_size)),
    A.CenterCrop(val_crop_size, val_crop_size),
    A.Normalize(mean=mean, std=std),
    ToTensor(),
])

train_loader, val_loader, train_eval_loader = get_train_val_loaders(
    data_path,
    train_transforms=train_transforms,
    val_transforms=val_transforms,
    batch_size=batch_size,
    num_workers=num_workers,
    val_batch_size=batch_size,
    pin_memory=True,
    train_sampler='distributed',
    val_sampler='distributed'
)

# Image denormalization function to plot predictions with images

def transforms_valid():
        return A.Compose(
            [
                A.Resize(256, 896),
                A.Normalize(mean=(Experiment.mean, Experiment.mean, Experiment.mean), std=(Experiment.std, Experiment.std, Experiment.std)),
                ChannelTranspose()
            ]

'tiles' - crop image in overlapping tiles (guaranteed to process entire dataset)
    :return: (train_loader, valid_loader)
    """

    if augmentation == "hard":
        train_transform = hard_augmentations()
    elif augmentation == "medium":
        train_transform = medium_augmentations()
    elif augmentation == "light":
        train_transform = light_augmentations()
    elif augmentation == "safe":
        train_transform = safe_augmentations()
    else:
        train_transform = A.Normalize()

    valid_transform = A.Normalize()
    assert train_mode in {"random", "tiles"}
    locations = TRAIN_LOCATIONS

    if train_mode == "random":

        train_data = []
        valid_data = []

        # For validation, we remove the first five images of every location (e.g., austin{1-5}.tif, chicago{1-5}.tif) from the training set.
        # That is suggested validation strategy by competition host

        if fast:
            # Fast training model. Use only one image per location for training and one image per location for validation
            for loc in locations:
                valid_data.append(f"{loc}1")
                train_data.append(f"{loc}6")

list_transforms.extend(
                [CropNonEmptyMaskIfExists(crop_image_size[0], crop_image_size[1], p=0.85),
                 HorizontalFlip(p=0.5),
                 VerticalFlip(p=0.5),
                 RandomBrightnessContrast(p=0.1, brightness_limit=0.1, contrast_limit=0.1)
                 ])
        else:
            list_transforms.extend(
                [HorizontalFlip(p=0.5),
                 VerticalFlip(p=0.5),
                 RandomBrightnessContrast(p=0.1, brightness_limit=0.1, contrast_limit=0.1)
                 ]
            )
    list_transforms.extend(
        [
            Normalize(mean=mean, std=std, p=1),
            ToTensor()
        ]
    )
    list_trfms = Compose(list_transforms)
    return list_trfms

def load(self, path="/model"):
        self.transform = Compose(
            [
                LongestMaxSize(max_size=224),
                PadIfNeeded(224, 224, border_mode=cv2.BORDER_CONSTANT),
                Normalize(),
                ToTensor(),
            ]
        )
        self.model = torch.jit.load(os.path.join(path, "model.pth"))
        with open(os.path.join(path, "tag2class.json")) as fin:
            self.tag2class = json.load(fin)
            self.class2tag = {v: k for k, v in self.tag2class.items()}
            logging.debug(f"class2tag: {self.class2tag}")