import argparse
import utils
from tifffile import imread
import numpy as np
import torch.backends.cudnn as cudnn
import ast
import torch
from torchvision import transforms, datasets
import os
from catalyst.data import DistributedSamplerWrapper
import sys
import vision_transformer as vits
from torchvision import models as torchvision_models
from torch import nn
import torch.distributed as dist
import re
import json


def extract_and_save_feature_pipeline(args):
    # ============ preparing image data ... ============

    torch.manual_seed(args.seed)
    
    if not args.images_are_RGB:
        #image loader compatible with multi-channel images
        #prepare selected channels and their corresponding mean and std
        selected_channels = list(map(int, args.selected_channels))
        transform = transforms.Compose([])
        if args.resize:
            transform.transforms.append(transforms.Resize(args.resize_length))
        if args.normalize:
            norm_per_channel = ast.literal_eval(args.norm_per_channel)
            mean_for_selected_channel, std_for_selected_channel = tuple([norm_per_channel[0][mean] for mean in selected_channels]), tuple([norm_per_channel[1][mean] for mean in selected_channels])
            print("normalize with mean: ", mean_for_selected_channel, " and std: ", std_for_selected_channel)
            transform.transforms.append(transforms.Normalize(mean=mean_for_selected_channel, std=std_for_selected_channel))

        def load_image(self, idx, args):
            path, target = self.samples[idx]
            image = imread(path)
            image = image[:, :, selected_channels]
            image = image.astype(float)
            if args.center_crop:
                image = torch.from_numpy(image).permute(2, 0, 1)
                transform = transforms.CenterCrop(args.center_crop)
                image = transform(image)
                image = image.permute(1, 2, 0)
                image = image.numpy()
            return image

        class Multichannel_dataset(datasets.ImageFolder):
                def __getitem__(self, idx):
                    path, target = self.samples[idx]
                    image_np = load_image(self, idx, args)
                    image_np = utils.normalize_numpy_0_to_1(image_np)
                    if utils.check_nan(image_np):
                        print("nan in image: ", path)
                        print('taking first image in dataset as replacement, so there are duplicates instead of nan values')
                        image_np = load_image(self, 0, args)
                        image_np = utils.normalize_numpy_0_to_1(image_np)
                    image = torch.from_numpy(image_np).permute(2, 0, 1)
                    if self.transform is not None:
                        image = self.transform(image)
                    if torch.isnan(image).any():
                        print("nan in image")
                    return image, idx
    
        dataset_total =  Multichannel_dataset(os.path.join(args.dataset_dir), transform=transform)
    
    else: #images are RGB
        transform = transforms.Compose([
        transforms.Resize(256, interpolation=3),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ])

        class ReturnIndexDataset(datasets.ImageFolder):
            def __getitem__(self, idx):
                img, lab = super(ReturnIndexDataset, self).__getitem__(idx)
                return img, idx
        dataset_total = ReturnIndexDataset(os.path.join(args.dataset_dir), transform=transform)

    # SAMPLER SECTION
    if  args.use_weighted_sampler and args.class_weights and args.num_samples:
            print('Using weighted sampler')
            num_samples = args.num_samples
            weights_per_class = ast.literal_eval(args.class_weights)
            weights_per_sample = [weights_per_class[dataset_total.samples[i][1]] for i in range(len(dataset_total.samples))]
            weighted_sampler = torch.utils.data.sampler.WeightedRandomSampler(weights_per_sample, args.num_samples, replacement=False)
            sampler = DistributedSamplerWrapper(weighted_sampler)
    
    elif args.scDINO_full_pipeline:
        validation_split = float(1-args.train_datasetsplit_fraction)
        shuffle_dataset = True
        dataset_size = len(dataset_total)
        indices = list(range(dataset_size))
        split = int(np.floor(validation_split * dataset_size))
        if shuffle_dataset :
            np.random.seed(args.seed)
            np.random.shuffle(indices)
        train_indices, val_indices = indices[split:], indices[:split]
        val_sampler = torch.utils.data.SubsetRandomSampler(val_indices)
        sampler = DistributedSamplerWrapper(val_sampler)
        num_samples = len(val_indices)
    
    elif args.test_datasetsplit_fraction!=1:
        validation_split = float(args.test_datasetsplit_fraction)
        shuffle_dataset = True
        dataset_size = len(dataset_total)
        indices = list(range(dataset_size))
        split = int(np.floor(validation_split * dataset_size))
        if shuffle_dataset :
            np.random.seed(args.seed)
            np.random.shuffle(indices)
        train_indices, val_indices = indices[split:], indices[:split]
        val_sampler = torch.utils.data.SubsetRandomSampler(val_indices)
        sampler = DistributedSamplerWrapper(val_sampler)
        num_samples = len(val_indices)

    else:
        print("Loading all images of the dataset")
        sampler = torch.utils.data.DistributedSampler(dataset_total, shuffle=False)
        num_samples = len(dataset_total)

    data_loader = torch.utils.data.DataLoader(
        dataset_total,
        sampler=sampler,
        batch_size=args.batch_size_per_gpu,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=False,
        shuffle=False)
    
    print("Data loader created")

    local_batch, local_lables = next(iter(data_loader))
    num_channels=local_batch.shape[1]

    print(f"Data loaded with {num_samples} images with a size of {local_batch.shape[2]}x{local_batch.shape[3]} with {num_channels} channels")

    # ============ building network ... ============
    if "vit" in args.arch:
        num_in_chans_pretrained = utils.get_pretrained_weights_in_chans(args.pretrained_weights)
        print(f"Pretrained weights have {num_in_chans_pretrained} input channels")
        model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0, in_chans=int(num_in_chans_pretrained))
        model.cuda()
        utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size)
        model.eval()
        print(f"Model {args.arch} {args.patch_size}x{args.patch_size} with {num_in_chans_pretrained} in_chans built.")
    else:
        print(f"Architecture {args.arch} non supported")
        sys.exit(1)


    # ============ adjusting ViT model option... ============

    def embedding_seq(custom_embedding_map, selected_channels):
            map_dict = ast.literal_eval(custom_embedding_map)
            map_dict = {int(k):int(v) for k,v in map_dict.items()}
            embedding_seq = [map_dict[int(input_channel)] for input_channel in selected_channels]
            return embedding_seq

    def build_weight_emb(embedding_seq, model):
        weights = model.patch_embed.proj.weight
        weights = weights[:,embedding_seq,:,:]
        model.patch_embed.proj.weight = nn.Parameter(weights)
        return model

    if not args.use_mean_patch_embedding and not args.use_custom_embedding_map:
        if num_channels != num_in_chans_pretrained:
            print(f"Error: Number of channels in the dataset ({num_channels}) and pretrained weights ({num_in_chans_pretrained}) are different")
            print(f"Use --use_mean_patch_embedding or --use_custom_embedding_map to adjust the number of channels")
            raise ValueError(f"Number of channels in the dataset ({num_channels}) and pretrained weights ({num_in_chans_pretrained}) are different")

    if not args.images_are_RGB:
        if args.use_mean_patch_embedding:
            average_conv2d_weights = torch.mean(model.patch_embed.proj.weight,1, keepdim=True)
            conv2d_weights_per_chan = average_conv2d_weights.repeat(1,num_channels,1,1)
            model.patch_embed.proj.weight = nn.Parameter(conv2d_weights_per_chan)
        elif args.use_custom_embedding_map:
            embedding_seq = embedding_seq(args.custom_embedding_map, args.selected_channels)
            model = build_weight_emb(embedding_seq, model)

    # ============ extract features ... ============
    print("Extracting features for train set...")
    features, index_all = extract_features(model, data_loader, args.use_cuda)

    if utils.get_rank() == 0:
        features = nn.functional.normalize(features, dim=1, p=2)
    
    image_names= [dataset_total.samples[i][0] for i in index_all]
    labels = [utils.fetch_foldername_of_img_location(dataset_total,i, args.folder_depth_for_labels) for i in index_all]
    return features, labels, image_names

@torch.no_grad()
def extract_features(model, data_loader, use_cuda=True, multiscale=False):
    metric_logger = utils.MetricLogger(delimiter="  ")
    features = None
    indices_all = []
    for samples, index in metric_logger.log_every(data_loader, 10):
        samples = samples.cuda(non_blocking=True)
        index = index.cuda(non_blocking=True)
        if multiscale:
            feats = utils.multi_scale(samples, model)
        else:
            feats = model(samples.float()).clone()

        # init storage feature matrix
        if dist.get_rank() == 0 and features is None:
            features = torch.zeros(0, feats.shape[-1])
            if use_cuda:
                features = features.cuda(non_blocking=True)
            print(f"Storing features into tensor of shape {features.shape}")

        # get indexes from all processes
        y_all = torch.empty(dist.get_world_size(), index.size(0), dtype=index.dtype, device=index.device)
        y_l = list(y_all.unbind(0))
        y_all_reduce = torch.distributed.all_gather(y_l, index, async_op=True)
        y_all_reduce.wait()
        index_all = torch.cat(y_l)

        # share features between processes
        feats_all = torch.empty(
            dist.get_world_size(),
            feats.size(0),
            feats.size(1),
            dtype=feats.dtype,
            device=feats.device,
        )
        output_l = list(feats_all.unbind(0))
        output_all_reduce = torch.distributed.all_gather(output_l, feats, async_op=True)
        output_all_reduce.wait()

        # update storage feature matrix
        if dist.get_rank() == 0:
            if use_cuda:
                features = torch.cat((features,torch.cat(output_l)),0)
                indices_all.extend(index_all.tolist())
            else:
                features = torch.cat((features.cpu(),torch.cat(output_l).cpu()),0)
                indices_all.extend(index_all.tolist())
    return features, indices_all

if __name__ == '__main__':
    parser = argparse.ArgumentParser('Computation of CLS features')
    #computation settings
    parser.add_argument('--name_of_run', default='/recent_run', type=str)
    parser.add_argument('--batch_size_per_gpu', default=30, type=int, help='Per-GPU batch-size')
    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
    parser.add_argument('--use_cuda', default=True, type=utils.bool_flag,
        help="Should we store the features on GPU? We recommend setting this to False if you encounter OOM")
    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
    parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
    parser.add_argument("--checkpoint_key", default="teacher", type=str,
        help='Key to use in the checkpoint (example: "teacher")')
    parser.add_argument('--num_workers', default=0, type=int, help='Number of data loading workers per GPU.')
    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
    #image dataset settings
    parser.add_argument('--dataset_dir', default='/path/to/imagenet/', type=str)
    parser.add_argument('--selected_channels', default=[0,1,2], nargs='+', help="""list of channel indexes of the .tiff images which should be used to create the tensors.""")
    parser.add_argument('--channel_dict', default=None, type=str,help="""name of the channels in format as dict channel_number, channel_name.""")
    parser.add_argument('--resize', default=False, type=utils.bool_flag, help="""if images should be resized""")
    parser.add_argument('--resize_length', default=None, help="""quadratic resize length to resize images""")
    parser.add_argument('--norm_per_channel', default="[(x, x, x, x, x), (x, x, x, x, x)]", type=str, help="""2x tuple of mean and std per channel typically values between 0 and 1""")
    parser.add_argument('--norm_per_channel_file', default=None, help="""path to file with mean and std per channel in json format.""")
    parser.add_argument('--center_crop', default=None, type=int, help="""center crop factor to crop images""")
    parser.add_argument('--normalize', default=False, type=utils.bool_flag, help="""normalize with mean and std per channel""")
    parser.add_argument('--patch_embedding_mapping', default=None, help="""change the patch embedding weights by inputting a string of the sequence of rearrangement of the model '[0,1,2]' or the string 'average_weights' or None""")
    parser.add_argument('--parse_params',help='Load settings from file in json format. Command line options override values in file.')
    parser.add_argument('--use_weighted_sampler', default=False, type=utils.bool_flag,help='Use weighted sampler for training.')
    parser.add_argument('--class_weights', default=None, help="""list of weights for each class""")
    parser.add_argument("--num_samples", default=None, type=int, help="Number of images to run in total.")
    parser.add_argument("--read_model_arch_dynamically", default=None, type=str, help="Read model architecture from pretrained weights")
    parser.add_argument("--use_mean_patch_embedding", default=False, type=utils.bool_flag,help="Use mean patch embedding instead of first patch embedding")
    parser.add_argument("--use_custom_embedding_map", default=False, type=utils.bool_flag, help="Use custom embedding map")
    parser.add_argument("--custom_embedding_map", default=None, type=dict, help="Custom embedding map")
    parser.add_argument("--scDINO_full_pipeline", default=False, type=utils.bool_flag, help="Using scDINO full pipeline")
    parser.add_argument('--full_ViT_name', default='full_vit_name', type=str, help='name channel combi ViT')
    parser.add_argument("--train_datasetsplit_fraction", default=0.8, type=float, help="when using scDINO full pipeline")
    parser.add_argument("--test_datasetsplit_fraction", default=0.8, type=float, help="when using downstream analysis only")
    parser.add_argument('--seed', default=42, type=int, help='Random seed.')
    parser.add_argument('--folder_depth_for_labels', default=0, type=int, help='Folder depth for labels. 0 means that the labels are the folder names where the images are stored. 1 means one level above and so on. e.g  path/to/images/labelwhen3/labelwhen2/labelwhen1/labelwhen0/image.tiff')

    #save settings
    parser.add_argument('--output_dir', default='.', type=str)

    # Added
    parser.add_argument('--images_are_RGB', default=False, type=utils.bool_flag)


    args = parser.parse_args()

    args.selected_channels = list(map(int, args.selected_channels[0].split(',')))
    args.channel_dict = dict(zip(args.selected_channels, args.channel_dict.split(',')))
    
    args.resize_length = (int(args.resize_length),int(args.resize_length))

    if args.parse_params:
        t_args = argparse.Namespace()
        t_args.__dict__.update(ast.literal_eval(args.parse_params))
        args = parser.parse_args(namespace=t_args)

        # read mean and std per channel
    if args.norm_per_channel_file:
        with open(args.norm_per_channel_file) as f:
            norm_per_channel_json = json.load(f)
            norm_per_channel = str([tuple(norm_per_channel_json['mean']), tuple(norm_per_channel_json['std'])])
            args.norm_per_channel = norm_per_channel

    #adjust model arch and patch size according to pretrained weights
    def adjust_model_architecture(args):
        model_name = args.pretrained_weights.split("/")[-1]
        number = re.findall(r'\d+', model_name)
        if number:
            args.patch_size = int(number[0])
        if "small" in model_name:
            args.arch = "vit_small"
        elif "base" in model_name:
            args.arch = "vit_base"

    if args.read_model_arch_dynamically:
        adjust_model_architecture(args)

    utils.init_distributed_mode(args)
    print("git:\n  {}\n".format(utils.get_sha()))
    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
    cudnn.benchmark = True

    #compute CLS features and save them
    features, labels, image_names = extract_and_save_feature_pipeline(args)

    # save features and labels
    if dist.get_rank() == 0:

        #translate channel indexes to channel names
        def get_channel_name_combi(channel_combi, channel_dict):
            name_of_channel_combi = ""
            for channel_number in iter(str(channel_combi)):
                name_of_channel_combi = "_".join([name_of_channel_combi, channel_dict[int(channel_number)]])
            return name_of_channel_combi

        #concatenate args.selected_channels to string
        selected_channel_str = "".join(str(x) for x in args.selected_channels)
        channel_names = get_channel_name_combi(selected_channel_str, args.channel_dict)

        if args.scDINO_full_pipeline:
            split_string = args.full_ViT_name.split("_epoch")
            ViT_name = split_string[0]
            epoch_num = split_string[1]
            cls_directory = os.path.join(args.output_dir,args.name_of_run, f"{ViT_name}_analyses/CLS_features")
            epoch_path = os.path.join(cls_directory, f"epoch{epoch_num}")
            print("epoch path: ", epoch_path)
            print("cls directory: ", cls_directory)
          
        else:
            cls_directory = os.path.join(args.output_dir,args.name_of_run,"CLS_features")
            if not os.path.exists(cls_directory):
                os.makedirs(cls_directory)
            model_name = args.pretrained_weights.split('/')[-1].split('.')[0]
            epoch_path = os.path.join(cls_directory, f"channel{channel_names}_model_{model_name}")
        
        np.savetxt(f"{epoch_path}_features.csv", features.cpu().numpy(), delimiter=",")
        with open(os.path.join(cls_directory, "run_log.txt"), "w") as f:
            f.write(f"computing features with seed {args.seed}: \n")
            f.write("parameters: \n")
            for arg in vars(args):
                f.write(f"{arg} : {getattr(args, arg)} \n")

3.test_cases/4.DDP/pyscripts/compute_CLS_features.py
aws-samples/awsome-distributed-training
import argparse
import utils
from tifffile import imread
import numpy as np
import torch.backends.cudnn as cudnn
import ast
import torch
from torchvision import transforms, datasets
import os
from catalyst.data import DistributedSamplerWrapper
import sys
import vision_transformer as vits
from torchvision import models as torchvision_models
from torch import nn
import torch.distributed as dist
import re
import json


def extract_and_save_feature_pipeline(args):
    # ============ preparing image data ... ============

    torch.manual_seed(args.seed)
    
    if not args.images_are_RGB:
        #image loader compatible with multi-channel images
        #prepare selected channels and their corresponding mean and std
        selected_channels = list(map(int, args.selected_channels))
        transform = transforms.Compose([])
        if args.resize:
            transform.transforms.append(transforms.Resize(args.resize_length))
        if args.normalize:
            norm_per_channel = ast.literal_eval(args.norm_per_channel)
            mean_for_selected_channel, std_for_selected_channel = tuple([norm_per_channel[0][mean] for mean in selected_channels]), tuple([norm_per_channel[1][mean] for mean in selected_channels])
            print("normalize with mean: ", mean_for_selected_channel, " and std: ", std_for_selected_channel)
            transform.transforms.append(transforms.Normalize(mean=mean_for_selected_channel, std=std_for_selected_channel))

        def load_image(self, idx, args):
            path, target = self.samples[idx]
            image = imread(path)
            image = image[:, :, selected_channels]
            image = image.astype(float)
            if args.center_crop:
                image = torch.from_numpy(image).permute(2, 0, 1)
                transform = transforms.CenterCrop(args.center_crop)
                image = transform(image)
                image = image.permute(1, 2, 0)
                image = image.numpy()
            return image

        class Multichannel_dataset(datasets.ImageFolder):
                def __getitem__(self, idx):
                    path, target = self.samples[idx]
                    image_np = load_image(self, idx, args)
                    image_np = utils.normalize_numpy_0_to_1(image_np)
                    if utils.check_nan(image_np):
                        print("nan in image: ", path)
                        print('taking first image in dataset as replacement, so there are duplicates instead of nan values')
                        image_np = load_image(self, 0, args)
                        image_np = utils.normalize_numpy_0_to_1(image_np)
                    image = torch.from_numpy(image_np).permute(2, 0, 1)
                    if self.transform is not None:
                        image = self.transform(image)
                    if torch.isnan(image).any():
                        print("nan in image")
                    return image, idx
    
        dataset_total =  Multichannel_dataset(os.path.join(args.dataset_dir), transform=transform)
    
    else: #images are RGB
        transform = transforms.Compose([
        transforms.Resize(256, interpolation=3),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ])

        class ReturnIndexDataset(datasets.ImageFolder):
            def __getitem__(self, idx):
                img, lab = super(ReturnIndexDataset, self).__getitem__(idx)
                return img, idx
        dataset_total = ReturnIndexDataset(os.path.join(args.dataset_dir), transform=transform)

    # SAMPLER SECTION
    if  args.use_weighted_sampler and args.class_weights and args.num_samples:
            print('Using weighted sampler')
            num_samples = args.num_samples
            weights_per_class = ast.literal_eval(args.class_weights)
            weights_per_sample = [weights_per_class[dataset_total.samples[i][1]] for i in range(len(dataset_total.samples))]
            weighted_sampler = torch.utils.data.sampler.WeightedRandomSampler(weights_per_sample, args.num_samples, replacement=False)
            sampler = DistributedSamplerWrapper(weighted_sampler)
    
    elif args.scDINO_full_pipeline:
        validation_split = float(1-args.train_datasetsplit_fraction)
        shuffle_dataset = True
        dataset_size = len(dataset_total)
        indices = list(range(dataset_size))
        split = int(np.floor(validation_split * dataset_size))
        if shuffle_dataset :
            np.random.seed(args.seed)
            np.random.shuffle(indices)
        train_indices, val_indices = indices[split:], indices[:split]
        val_sampler = torch.utils.data.SubsetRandomSampler(val_indices)
        sampler = DistributedSamplerWrapper(val_sampler)
        num_samples = len(val_indices)
    
    elif args.test_datasetsplit_fraction!=1:
        validation_split = float(args.test_datasetsplit_fraction)
        shuffle_dataset = True
        dataset_size = len(dataset_total)
        indices = list(range(dataset_size))
        split = int(np.floor(validation_split * dataset_size))
        if shuffle_dataset :
            np.random.seed(args.seed)
            np.random.shuffle(indices)
        train_indices, val_indices = indices[split:], indices[:split]
        val_sampler = torch.utils.data.SubsetRandomSampler(val_indices)
        sampler = DistributedSamplerWrapper(val_sampler)
        num_samples = len(val_indices)

    else:
        print("Loading all images of the dataset")
        sampler = torch.utils.data.DistributedSampler(dataset_total, shuffle=False)
        num_samples = len(dataset_total)

    data_loader = torch.utils.data.DataLoader(
        dataset_total,
        sampler=sampler,
        batch_size=args.batch_size_per_gpu,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=False,
        shuffle=False)
    
    print("Data loader created")

    local_batch, local_lables = next(iter(data_loader))
    num_channels=local_batch.shape[1]

    print(f"Data loaded with {num_samples} images with a size of {local_batch.shape[2]}x{local_batch.shape[3]} with {num_channels} channels")

    # ============ building network ... ============
    if "vit" in args.arch:
        num_in_chans_pretrained = utils.get_pretrained_weights_in_chans(args.pretrained_weights)
        print(f"Pretrained weights have {num_in_chans_pretrained} input channels")
        model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0, in_chans=int(num_in_chans_pretrained))
        model.cuda()
        utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size)
        model.eval()
        print(f"Model {args.arch} {args.patch_size}x{args.patch_size} with {num_in_chans_pretrained} in_chans built.")
    else:
        print(f"Architecture {args.arch} non supported")
        sys.exit(1)


    # ============ adjusting ViT model option... ============

    def embedding_seq(custom_embedding_map, selected_channels):
            map_dict = ast.literal_eval(custom_embedding_map)
            map_dict = {int(k):int(v) for k,v in map_dict.items()}
            embedding_seq = [map_dict[int(input_channel)] for input_channel in selected_channels]
            return embedding_seq

    def build_weight_emb(embedding_seq, model):
        weights = model.patch_embed.proj.weight
        weights = weights[:,embedding_seq,:,:]
        model.patch_embed.proj.weight = nn.Parameter(weights)
        return model

    if not args.use_mean_patch_embedding and not args.use_custom_embedding_map:
        if num_channels != num_in_chans_pretrained:
            print(f"Error: Number of channels in the dataset ({num_channels}) and pretrained weights ({num_in_chans_pretrained}) are different")
            print(f"Use --use_mean_patch_embedding or --use_custom_embedding_map to adjust the number of channels")
            raise ValueError(f"Number of channels in the dataset ({num_channels}) and pretrained weights ({num_in_chans_pretrained}) are different")

    if not args.images_are_RGB:
        if args.use_mean_patch_embedding:
            average_conv2d_weights = torch.mean(model.patch_embed.proj.weight,1, keepdim=True)
            conv2d_weights_per_chan = average_conv2d_weights.repeat(1,num_channels,1,1)
            model.patch_embed.proj.weight = nn.Parameter(conv2d_weights_per_chan)
        elif args.use_custom_embedding_map:
            embedding_seq = embedding_seq(args.custom_embedding_map, args.selected_channels)
            model = build_weight_emb(embedding_seq, model)

    # ============ extract features ... ============
    print("Extracting features for train set...")
    features, index_all = extract_features(model, data_loader, args.use_cuda)

    if utils.get_rank() == 0:
        features = nn.functional.normalize(features, dim=1, p=2)
    
    image_names= [dataset_total.samples[i][0] for i in index_all]
    labels = [utils.fetch_foldername_of_img_location(dataset_total,i, args.folder_depth_for_labels) for i in index_all]
    return features, labels, image_names

@torch.no_grad()
def extract_features(model, data_loader, use_cuda=True, multiscale=False):
    metric_logger = utils.MetricLogger(delimiter="  ")
    features = None
    indices_all = []
    for samples, index in metric_logger.log_every(data_loader, 10):
        samples = samples.cuda(non_blocking=True)
        index = index.cuda(non_blocking=True)
        if multiscale:
            feats = utils.multi_scale(samples, model)
        else:
            feats = model(samples.float()).clone()

        # init storage feature matrix
        if dist.get_rank() == 0 and features is None:
            features = torch.zeros(0, feats.shape[-1])
            if use_cuda:
                features = features.cuda(non_blocking=True)
            print(f"Storing features into tensor of shape {features.shape}")

        # get indexes from all processes
        y_all = torch.empty(dist.get_world_size(), index.size(0), dtype=index.dtype, device=index.device)
        y_l = list(y_all.unbind(0))
        y_all_reduce = torch.distributed.all_gather(y_l, index, async_op=True)
        y_all_reduce.wait()
        index_all = torch.cat(y_l)

        # share features between processes
        feats_all = torch.empty(
            dist.get_world_size(),
            feats.size(0),
            feats.size(1),
            dtype=feats.dtype,
            device=feats.device,
        )
        output_l = list(feats_all.unbind(0))
        output_all_reduce = torch.distributed.all_gather(output_l, feats, async_op=True)
        output_all_reduce.wait()

        # update storage feature matrix
        if dist.get_rank() == 0:
            if use_cuda:
                features = torch.cat((features,torch.cat(output_l)),0)
                indices_all.extend(index_all.tolist())
            else:
                features = torch.cat((features.cpu(),torch.cat(output_l).cpu()),0)
                indices_all.extend(index_all.tolist())
    return features, indices_all

if __name__ == '__main__':
    parser = argparse.ArgumentParser('Computation of CLS features')
    #computation settings
    parser.add_argument('--name_of_run', default='/recent_run', type=str)
    parser.add_argument('--batch_size_per_gpu', default=30, type=int, help='Per-GPU batch-size')
    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
    parser.add_argument('--use_cuda', default=True, type=utils.bool_flag,
        help="Should we store the features on GPU? We recommend setting this to False if you encounter OOM")
    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
    parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
    parser.add_argument("--checkpoint_key", default="teacher", type=str,
        help='Key to use in the checkpoint (example: "teacher")')
    parser.add_argument('--num_workers', default=0, type=int, help='Number of data loading workers per GPU.')
    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
    #image dataset settings
    parser.add_argument('--dataset_dir', default='/path/to/imagenet/', type=str)
    parser.add_argument ('--images_are_RGB',help='If images are RGB, set this to True. If images are grayscale, set this to False.', default=False, type=utils.bool_flag)
    parser.add_argument('--selected_channels', default=[0,1,2], nargs='+', help="""list of channel indexes of the .tiff images which should be used to create the tensors.""")
    parser.add_argument('--channel_dict', default=None, type=str,help="""name of the channels in format as dict channel_number, channel_name.""")
    parser.add_argument('--resize', default=False, help="""if images should be resized""")
    parser.add_argument('--resize_length', default=None, help="""quadratic resize length to resize images""")
    parser.add_argument('--norm_per_channel', default="[(x, x, x, x, x), (x, x, x, x, x)]", type=str, help="""2x tuple of mean and std per channel typically values between 0 and 1""")
    parser.add_argument('--norm_per_channel_file', default=None, help="""path to file with mean and std per channel in json format.""")
    parser.add_argument('--center_crop', type=int, default=None, help="""center crop factor to crop images""")
    parser.add_argument('--normalize', default="False", type=str, help="""normalize with mean and std per channel""")
    parser.add_argument('--patch_embedding_mapping', default=None, help="""change the patch embedding weights by inputting a string of the sequence of rearrangement of the model '[0,1,2]' or the string 'average_weights' or None""")
    parser.add_argument('--parse_params',help='Load settings from file in json format. Command line options override values in file.')
    parser.add_argument('--use_weighted_sampler', default=False, type=bool, help='Use weighted sampler for training.')
    parser.add_argument('--class_weights', default=None, help="""list of weights for each class""")
    parser.add_argument("--num_samples", default=None, type=int, help="Number of images to run in total.")
    parser.add_argument("--read_model_arch_dynamically", default=None, type=str, help="Read model architecture from pretrained weights")
    parser.add_argument("--use_mean_patch_embedding", default=False, type=bool, help="Use mean patch embedding instead of first patch embedding")
    parser.add_argument("--use_custom_embedding_map", default=False, type=bool, help="Use custom embedding map")
    parser.add_argument("--custom_embedding_map", default=None, type=dict, help="Custom embedding map")
    parser.add_argument("--scDINO_full_pipeline", default=False, type=bool, help="Using scDINO full pipeline")
    parser.add_argument('--full_ViT_name', default='full_vit_name', type=str, help='name channel combi ViT')
    parser.add_argument("--train_datasetsplit_fraction", default=0.8, type=float, help="when using scDINO full pipeline")
    parser.add_argument("--test_datasetsplit_fraction", default=0.8, type=float, help="when using downstream analysis only")
    parser.add_argument('--seed', default=42, type=int, help='Random seed.')
    parser.add_argument('--folder_depth_for_labels', default=0, type=int, help='Folder depth for labels. 0 means that the labels are the folder names where the images are stored. 1 means one level above and so on. e.g  path/to/images/labelwhen3/labelwhen2/labelwhen1/labelwhen0/image.tiff')

    #save settings
    parser.add_argument('--output_dir', default='.', type=str)
    args = parser.parse_args()

    args.selected_channels = list(map(int, args.selected_channels[0].split(',')))
    args.channel_dict = dict(zip(args.selected_channels, args.channel_dict.split(',')))
    args.resize_length = (int(args.resize_length),int(args.resize_length))

    if args.parse_params:
        t_args = argparse.Namespace()
        t_args.__dict__.update(ast.literal_eval(args.parse_params))
        args = parser.parse_args(namespace=t_args)

        # read mean and std per channel
    if args.norm_per_channel_file:
        with open(args.norm_per_channel_file) as f:
            norm_per_channel_json = json.load(f)
            norm_per_channel = str([tuple(norm_per_channel_json['mean']), tuple(norm_per_channel_json['std'])])
            args.norm_per_channel = norm_per_channel

    #adjust model arch and patch size according to pretrained weights
    def adjust_model_architecture(args):
        model_name = args.pretrained_weights.split("/")[-1]
        number = re.findall(r'\d+', model_name)
        if number:
            args.patch_size = int(number[0])
        if "small" in model_name:
            args.arch = "vit_small"
        elif "base" in model_name:
            args.arch = "vit_base"

    if args.read_model_arch_dynamically:
        adjust_model_architecture(args)

    utils.init_distributed_mode(args)
    print("git:\n  {}\n".format(utils.get_sha()))
    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
    cudnn.benchmark = True

    #compute CLS features and save them
    features, labels, image_names = extract_and_save_feature_pipeline(args)
    # save features and labels as files
    if dist.get_rank() == 0:

        #translate channel indexes to channel names
        def get_channel_name_combi(channel_combi, channel_dict):
            name_of_channel_combi = ""
            for channel_number in iter(str(channel_combi)):
                name_of_channel_combi = "_".join([name_of_channel_combi, channel_dict[int(channel_number)]])
            return name_of_channel_combi

        #concatenate args.selected_channels to string
        selected_channel_str = "".join(str(x) for x in args.selected_channels)
        channel_names = get_channel_name_combi(selected_channel_str, args.channel_dict)

        if args.scDINO_full_pipeline:
            path = os.path.join(args.output_dir,args.name_of_run)+"/"
          
        else: 
             path = os.path.join(args.output_dir,args.name_of_run, f"CLS_features/")
            
        np.savetxt(f"{path}class_labels.csv", labels, delimiter=",", fmt="%s")
        np.savetxt(f"{path}image_paths.csv", image_names, delimiter=",", fmt="%s")

3.test_cases/4.DDP/pyscripts/extract_image_labels.py
aws-samples/awsome-distributed-training
import pandas as pd
import torch
from sklearn import preprocessing
import numpy as np
import yaml
import os
import glob

def get_channel_name_combi(channel_combi_num, channel_dict):
    name_of_channel_combi = ""
    for channel_number in iter(str(channel_combi_num)):
        name_of_channel_combi = "_".join([name_of_channel_combi, channel_dict[int(channel_number)]])
    return name_of_channel_combi

def get_channel_number_combi(channel_names, channel_dict):
    channel_combi = ""
    for channel_name in channel_names.split('_'):
        for key, value in channel_dict.items():
            if value == channel_name:
                channel_combi = "".join([channel_combi, str(key)])
    return channel_combi

def get_channel_name_combi_list(selected_channels, channel_dict):
    channel_names = []
    for channel_combi in selected_channels:
        channel_names.append(get_channel_name_combi(channel_combi,channel_dict))
    return channel_names

config_file = '../scDINO_full_pipeline.yaml'
with open(config_file, "r") as f:
    config = yaml.load(f,Loader=yaml.FullLoader)

name_of_run = config['meta']['name_of_run']
sk_save_dir = config['meta']['output_dir']
save_dir_downstream_run = sk_save_dir+"/"+name_of_run


#channel_aTub_model_checkpoint2_features.csv
features_path = f"{save_dir_downstream_run}/CLS_features/"

features_file = glob.glob(features_path+'*_features.csv')[0]

labels_file = f"{save_dir_downstream_run}/CLS_features/class_labels.csv"


dino_vit_name = config['train_scDINO']['dino_vit_name']

selected_channels = [config['meta']['selected_channel_combination_per_run']]
selected_channels = list(eval(selected_channels[0]))

channel_dict = config['meta']['channel_dict']
channel_dict = dict(zip(selected_channels, channel_dict.split(',')))
channel_names=get_channel_name_combi_list(selected_channels, channel_dict)

save_dir = f"{save_dir_downstream_run}/{dino_vit_name}_channel{channel_names[0]}_analyses/"

###iterate through all runs and save the results in one file
##for run_index, features in enumerate(snakemake.input["features"]):
    #data

features = np.genfromtxt(features_file, delimiter = ',')
class_labels_pd = pd.read_csv(labels_file, header=None)
class_labels = class_labels_pd[0].tolist()
#params
run_name= name_of_run

# if snakemake.params['scDINO_full_pipeline']:
#     save_dir = snakemake.params['save_dir'][0]
#     neighbors = snakemake.config['downstream_analyses']['kNN']['global']['n_neighbors']
#     temperature = snakemake.config['downstream_analyses']['kNN']['global']['temperature']
# else:

save_dir = save_dir
neighbors = config['downstream_analyses']['kNN']['global']['n_neighbors']
print('NEIGHBORS = ')
print(neighbors)

temperature = config['downstream_analyses']['kNN']['global']['temperature']
print(temperature)

def number_of_classes(labels):
    return len(set(labels))

def label_to_number(labels):
    le = preprocessing.LabelEncoder()
    le.fit(labels)
    labels_numbers = le.transform(labels)
    return labels_numbers

indices = np.arange(features.shape[0])
np.random.seed(config['meta']['seed'])
np.random.shuffle(indices)
features = features[indices]
class_labels = np.array(class_labels)[indices]
train_features = torch.from_numpy(features[:int(features.shape[0]*0.8),:])
train_labels = torch.from_numpy(label_to_number(class_labels[:int(features.shape[0]*0.8)]))
test_features = torch.from_numpy(features[int(features.shape[0]*0.8):,:])
test_labels = torch.from_numpy(label_to_number(class_labels[int(features.shape[0]*0.8):]))
num_classes = number_of_classes(class_labels)

@torch.no_grad()
def knn_classifier(train_features, train_labels, test_features, test_labels, k, T, num_classes):
    top1, top2, total = 0.0, 0.0, 0
    train_features = train_features.t()
    num_test_images, num_chunks = test_labels.shape[0], 100
    imgs_per_chunk = num_test_images // num_chunks
    retrieval_one_hot = torch.zeros(k, num_classes).to(train_features.device)
    for idx in range(0, num_test_images, imgs_per_chunk):
        # get the features for test images
        features = test_features[
            idx : min((idx + imgs_per_chunk), num_test_images), :
        ]
        targets = test_labels[idx : min((idx + imgs_per_chunk), num_test_images)]
        batch_size = targets.shape[0]

        # calculate the dot product and compute top-k neighbors
        similarity = torch.mm(features, train_features)
        distances, indices = similarity.topk(k, largest=True, sorted=True)
        candidates = train_labels.view(1, -1).expand(batch_size, -1)
        retrieved_neighbors = torch.gather(candidates, 1, indices)

        retrieval_one_hot.resize_(batch_size * k, num_classes).zero_()
        retrieval_one_hot.scatter_(1, retrieved_neighbors.view(-1, 1), 1)
        distances_transform = distances.clone().div_(T).exp_()
        probs = torch.sum(
            torch.mul(
                retrieval_one_hot.view(batch_size, -1, num_classes),
                distances_transform.view(batch_size, -1, 1),
            ),
            1,
        )
        _, predictions = probs.sort(1, True)

        # find the predictions that match the target
        correct = predictions.eq(targets.data.view(-1, 1))
        top1 = top1 + correct.narrow(1, 0, 1).sum().item()
        top2 = top2 + correct.narrow(1, 0, min(2, k)).sum().item()  #dependent on number of classes
        total += targets.size(0)
        top1 = top1 * 100.0 / total
        top2 = top2 * 100.0 / total
        return top1, top2

    
knn_output_dir = save_dir+"/kNN/"

if not os.path.exists(knn_output_dir):
    os.makedirs(knn_output_dir)
with open(knn_output_dir+"global_kNN.txt", 'a') as f:
    f.write("Run: "+run_name+"\n")
    for k in neighbors:
        try:
            top1, top2 = knn_classifier(train_features, train_labels,test_features, test_labels, k, temperature, num_classes=num_classes)
            f.write(f"{k}-NN Top1: {round(top1)}, Top2: {round(top2)}\n")
            #delete variables
            del top1, top2
        except Exception as e:
            f.write("Error: "+str(e)+"\n")             
    f.write("--------------------------------------------------\n")


3.test_cases/4.DDP/pyscripts/global_kNN.py
aws-samples/awsome-distributed-training
import argparse
import os
from random import shuffle
import sys
import datetime
import time
import math
import json
from pathlib import Path
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torchvision import models as torchvision_models
from vision_transformer import DINOHead
import utils
import vision_transformer as vits
from tifffile import imread
import numpy
from PIL import Image
import torchvision
import ast
from catalyst.data import DistributedSamplerWrapper

torchvision_archs = sorted(name for name in torchvision_models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(torchvision_models.__dict__[name]))

def get_args_parser():
    parser = argparse.ArgumentParser('DINO', add_help=False)

    # Model parameters
    parser.add_argument('--arch', default='vit_small', type=str,
        choices=['vit_tiny', 'vit_small', 'vit_base', 'xcit', 'deit_tiny', 'deit_small'] \
                + torchvision_archs,
        help="""Name of architecture to train. For quick experiments with ViTs,
        we recommend using vit_tiny or vit_small.""")
    parser.add_argument('--patch_size', default=16, type=int, help="""Size in pixels
        of input square patches - default 16 (for 16x16 patches). Using smaller
        values leads to better performance but requires more memory. Applies only
        for ViTs (vit_tiny, vit_small and vit_base). If <16, we recommend disabling
        mixed precision training (--use_fp16 false) to avoid unstabilities.""")
    parser.add_argument('--out_dim', default=65536, type=int, help="""Dimensionality of
        the DINO head output. For complex and large datasets large values (like 65k) work well.""")
    parser.add_argument('--norm_last_layer', default=True, type=utils.bool_flag,
        help="""Whether or not to weight normalize the last layer of the DINO head.
        Not normalizing leads to better performance but can make the training unstable.
        In our experiments, we typically set this paramater to False with vit_small and True with vit_base.""")
    parser.add_argument('--momentum_teacher', default=0.996, type=float, help="""Base EMA
        parameter for teacher update. The value is increased to 1 during training with cosine schedule.
        We recommend setting a higher value with small batches: for example use 0.9995 with batch size of 256.""")
    parser.add_argument('--use_bn_in_head', default=False, type=utils.bool_flag,
        help="Whether to use batch normalizations in projection head (Default: False)")

    # Temperature teacher parameters
    parser.add_argument('--warmup_teacher_temp', default=0.04, type=float,
        help="""Initial value for the teacher temperature: 0.04 works well in most cases.
        Try decreasing it if the training loss does not decrease.""")
    parser.add_argument('--teacher_temp', default=0.04, type=float, help="""Final value (after linear warmup)
        of the teacher temperature. For most experiments, anything above 0.07 is unstable. We recommend
        starting with the default value of 0.04 and increase this slightly if needed.""")
    parser.add_argument('--warmup_teacher_temp_epochs', default=0, type=int,
        help='Number of warmup epochs for the teacher temperature (Default: 30).')

    # Training/Optimization parameters
    parser.add_argument('--use_fp16', type=utils.bool_flag, default=True, help="""Whether or not
        to use half precision for training. Improves training time and memory requirements,
        but can provoke instability and slight decay of performance. We recommend disabling
        mixed precision if the loss is unstable, if reducing the patch size or if training with bigger ViTs.""")
    parser.add_argument('--weight_decay', type=float, default=0.04, help="""Initial value of the
        weight decay. With ViT, a smaller value at the beginning of training works well.""")
    parser.add_argument('--weight_decay_end', type=float, default=0.4, help="""Final value of the
        weight decay. We use a cosine schedule for WD and using a larger decay by
        the end of training improves performance for ViTs.""")
    parser.add_argument('--clip_grad', type=float, default=3.0, help="""Maximal parameter
        gradient norm if using gradient clipping. Clipping with norm .3 ~ 1.0 can
        help optimization for larger ViT architectures. 0 for disabling.""")
    parser.add_argument('--batch_size_per_gpu', default=64, type=int,
        help='Per-GPU batch-size : number of distinct images loaded on one GPU.')
    parser.add_argument('--epochs', default=100, type=int, help='Number of epochs of training.')
    parser.add_argument('--freeze_last_layer', default=1, type=int, help="""Number of epochs
        during which we keep the output layer fixed. Typically doing so during
        the first epoch helps training. Try increasing this value if the loss does not decrease.""")
    parser.add_argument("--lr", default=0.0005, type=float, help="""Learning rate at the end of
        linear warmup (highest LR used during training). The learning rate is linearly scaled
        with the batch size, and specified here for a reference batch size of 256.""")
    parser.add_argument("--warmup_epochs", default=10, type=int,
        help="Number of epochs for the linear learning-rate warm up.")
    parser.add_argument('--min_lr', type=float, default=1e-6, help="""Target LR at the
        end of optimization. We use a cosine LR schedule with linear warmup.""")
    parser.add_argument('--optimizer', default='adamw', type=str,
        choices=['adamw', 'sgd', 'lars'], help="""Type of optimizer. We recommend using adamw with ViTs.""")
    parser.add_argument('--drop_path_rate', type=float, default=0.1, help="stochastic depth rate")
    parser.add_argument('--selected_channels', default=[0,1,2], nargs='+', help="""list of channel indexes of the .tiff images which should be used to create the tensors.""")
    parser.add_argument('--norm_per_channel_file', default="path/to/file", type=str, help="""path to mean and std file in json format""")
    parser.add_argument('--upscale_factor', default=0, type=float, help="""upscale factor to upsample images""")
    parser.add_argument('--center_crop', default=0, type=int, help="""center crop factor to crop images""")

    # Multi-crop parameters
    parser.add_argument('--global_crops_scale', type=float, nargs='+', default=(0.4, 1.),
        help="""Scale range of the cropped image before resizing, relatively to the origin image.
        Used for large global view cropping. When disabling multi-crop (--local_crops_number 0), we
        recommand using a wider range of scale ("--global_crops_scale 0.14 1." for example)""")
    parser.add_argument('--local_crops_number', type=int, default=8, help="""Number of small
        local views to generate. Set this parameter to 0 to disable multi-crop training.
        When disabling multi-crop we recommend to use "--global_crops_scale 0.14 1." """)
    parser.add_argument('--local_crops_scale', type=float, nargs='+', default=(0.05, 0.4),
        help="""Scale range of the cropped image before resizing, relatively to the origin image.
        Used for small local view cropping of multi-crop.""")

    # Misc
    parser.add_argument('--output_dir', default=".", type=str, help='Path to save logs and checkpoints.')
    parser.add_argument('--saveckp_freq', default=5, type=int, help='Save checkpoint every x epochs.')
    parser.add_argument('--seed', default=42, type=int, help='Random seed.')
    parser.add_argument('--num_workers', default=10, type=int, help='Number of data loading workers per GPU.')
    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
    parser.add_argument('--parse_params',help='Load settings from file in json format. Command line options override values in file.')
    parser.add_argument('--train_datasetsplit_fraction', default=0.8, type=float, help='fraction of training dataset to use for training')
    parser.add_argument('--dino_vit_name', default='vit_small_p16', type=str, help='name of ViTs')
    parser.add_argument('--channel_dict', default=None, type=str,help="""name of the channels in format as dict channel_number, channel_name.""")
    parser.add_argument('--name_of_run', default='/recent_run', type=str)
    parser.add_argument('--full_ViT_name', default='full_vit_name', type=str, help='name channel combi ViT')
    
    # Added
    parser.add_argument('--dataset_dir', default=".")
    parser.add_argument('--images_are_RGB', default=False, action="store_true")


    return parser


def train_dino(args, save_dir):
    utils.init_distributed_mode(args)
    utils.fix_random_seeds(args.seed)
    print("git:\n  {}\n".format(utils.get_sha()))
    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
    cudnn.benchmark = True
    print('saving log file of run parameters')
    with open(os.path.join(save_dir, "run_log.txt"), "w") as f:
        f.write(f"successfully computed features with seed {args.seed}: \n")
        f.write("parameters: \n")
        for arg in vars(args):
            f.write(f"{arg} : {getattr(args, arg)} \n")

    # ============ preparing data ... ============

    def load_mean_std_per_channel(norm_per_channel_file):
        with open(args.norm_per_channel_file) as f:
            norm_per_channel_json = json.load(f)
            norm_per_channel = [norm_per_channel_json['mean'], norm_per_channel_json['std']]
        return norm_per_channel

    def create_mean_std_per_channel_for_channel_combi(norm_per_channel_file, selected_channels):
        norm_per_channel = load_mean_std_per_channel(args.norm_per_channel_file)
        mean_for_selected_channel, std_for_selected_channel = tuple([norm_per_channel[0][mean] for mean in selected_channels]), tuple([norm_per_channel[1][std] for std in selected_channels])
        return mean_for_selected_channel, std_for_selected_channel
    
    selected_channels = list(map(int, args.selected_channels))
    mean_for_selected_channel, std_for_selected_channel = create_mean_std_per_channel_for_channel_combi(args.norm_per_channel_file, selected_channels)
    
    transform = DataAugmentationDINO(
        args.images_are_RGB,
        args.global_crops_scale,
        args.local_crops_scale,
        args.local_crops_number,
        mean_for_selected_channel, 
        std_for_selected_channel,
    )

    if not args.images_are_RGB:

        print('ENTERING CUSTOM DATSET CLASS')

        class Multichannel_dataset(datasets.ImageFolder):
                def __getitem__(self, idx):
                    path, target = self.samples[idx]
                    image_np= imread(path)
                    image_np=image_np.astype(float)
                    image_np = image_np[:,:,selected_channels]
                    if args.center_crop:
                        image = torch.from_numpy(image_np).permute(2, 0, 1)
                        transform = transforms.CenterCrop(args.center_crop)
                        image = transform(image)
                        image = image.permute(1, 2, 0)
                        image_np = image.detach().cpu().numpy()
                    image_np = utils.normalize_numpy_0_to_1(image_np)
                    if utils.check_nan(image_np):
                        print("nan in image: ", path)
                        return None
                    else:
                        image = torch.from_numpy(image_np).float().permute(2, 0, 1)
                        if self.transform is not None:
                            image = self.transform(image)
                        if self.target_transform is not None:
                            target = self.target_transform(target)
                        return image, idx

        dataset_total =  Multichannel_dataset(os.path.join(args.dataset_dir), transform=transform)
    else:
        dataset_total = datasets.ImageFolder(args.dataset_dir, transform=transform)

    #SAMPLER SECTION
    validation_split = float(1-args.train_datasetsplit_fraction)
    shuffle_dataset = True
    # Creating data indices for training and validation splits:
    dataset_size = len(dataset_total)
    indices = list(range(dataset_size))
    split = int(numpy.floor(validation_split * dataset_size))
    if shuffle_dataset :
        numpy.random.seed(args.seed)
        numpy.random.shuffle(indices)
    train_indices, val_indices = indices[split:], indices[:split]

    print("Split between train and test dataset:",len(train_indices), len(val_indices))

    print(f"Train dataset consists of {len(train_indices)} images.")

    # Creating data samplers and loaders:
    train_sampler = torch.utils.data.SubsetRandomSampler(train_indices)
    train_sampler_wrapped = DistributedSamplerWrapper(train_sampler)

    data_loader = torch.utils.data.DataLoader(
        dataset_total,
        sampler=train_sampler_wrapped,
        batch_size=args.batch_size_per_gpu,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=False,
        shuffle=False, 
        collate_fn=utils.collate_fn)

    print("Successfully loaded data.")

    # ============ building student and teacher networks ... ============
    # we changed the name DeiT-S for ViT-S to avoid confusions
    args.arch = args.arch.replace("deit", "vit")
    # if the network is a Vision Transformer (i.e. vit_tiny, vit_small, vit_base)
    if args.arch in vits.__dict__.keys():
        student = vits.__dict__[args.arch](
            patch_size=args.patch_size,
            drop_path_rate=args.drop_path_rate,
            in_chans= len(selected_channels)
        )
        teacher = vits.__dict__[args.arch](patch_size=args.patch_size, in_chans= len(selected_channels))
        embed_dim = student.embed_dim
    # otherwise, we check if the architecture is in torchvision models
    elif args.arch in torchvision_models.__dict__.keys():
        student = torchvision_models.__dict__[args.arch]()
        teacher = torchvision_models.__dict__[args.arch]()
        embed_dim = student.fc.weight.shape[1]
    else:
        print(f"Unknown architecture: {args.arch}")

    # multi-crop wrapper handles forward with inputs of different resolutions
    student = utils.MultiCropWrapper(student, DINOHead(
        embed_dim,
        args.out_dim,
        use_bn=args.use_bn_in_head,
        norm_last_layer=args.norm_last_layer,
    ))
    teacher = utils.MultiCropWrapper(
        teacher,
        DINOHead(embed_dim, args.out_dim, args.use_bn_in_head),
    )
    # move networks to gpu
    student, teacher = student.cuda(), teacher.cuda()
    # synchronize batch norms (if any)
    if utils.has_batchnorms(student):
        student = nn.SyncBatchNorm.convert_sync_batchnorm(student)
        teacher = nn.SyncBatchNorm.convert_sync_batchnorm(teacher)

        # we need DDP wrapper to have synchro batch norms working...
        teacher = nn.parallel.DistributedDataParallel(teacher, device_ids=[args.gpu])
        teacher_without_ddp = teacher.module
    else:
        # teacher_without_ddp and teacher are the same thing
        teacher_without_ddp = teacher
    student = nn.parallel.DistributedDataParallel(student, device_ids=[args.gpu])
    # teacher and student start with the same weights
    teacher_without_ddp.load_state_dict(student.module.state_dict())
    # there is no backpropagation through the teacher, so no need for gradients
    for p in teacher.parameters():
        p.requires_grad = False
    print(f"Student and Teacher are built: they are both {args.arch} network.")

    # ============ preparing loss ... ============
    dino_loss = DINOLoss(
        args.out_dim,
        args.local_crops_number + 2,  # total number of crops = 2 global crops + local_crops_number
        args.warmup_teacher_temp,
        args.teacher_temp,
        args.warmup_teacher_temp_epochs,
        args.epochs,
    ).cuda()

    # ============ preparing optimizer ... ============
    params_groups = utils.get_params_groups(student)
    if args.optimizer == "adamw":
        optimizer = torch.optim.AdamW(params_groups)  # to use with ViTs
    elif args.optimizer == "sgd":
        optimizer = torch.optim.SGD(params_groups, lr=0, momentum=0.9)  # lr is set by scheduler
    elif args.optimizer == "lars":
        optimizer = utils.LARS(params_groups)  # to use with convnet and large batches
    # for mixed precision training
    fp16_scaler = None
    if args.use_fp16:
        fp16_scaler = torch.cuda.amp.GradScaler()

    # ============ init schedulers ... ============
    lr_schedule = utils.cosine_scheduler(
        args.lr * (args.batch_size_per_gpu * utils.get_world_size()) / 256.,  # linear scaling rule
        args.min_lr,
        args.epochs, len(data_loader),
        warmup_epochs=args.warmup_epochs,
    )
    wd_schedule = utils.cosine_scheduler(
        args.weight_decay,
        args.weight_decay_end,
        args.epochs, len(data_loader),
    )
    # momentum parameter is increased to 1. during training with a cosine schedule
    momentum_schedule = utils.cosine_scheduler(args.momentum_teacher, 1,
                                               args.epochs, len(data_loader))
    print(f"Loss, optimizer and schedulers ready.")

    # ============ optionally resume training ... ============
    to_restore = {"epoch": 0}
    utils.restart_from_checkpoint(
        os.path.join(save_dir, "checkpoint.pth"),
        run_variables=to_restore,
        student=student,
        teacher=teacher,
        optimizer=optimizer,
        fp16_scaler=fp16_scaler,
        dino_loss=dino_loss,
    )
    start_epoch = to_restore["epoch"]

    start_time = time.time()
    print("Starting DINO training !")
    for epoch in range(start_epoch, args.epochs):
        data_loader.sampler.set_epoch(epoch)

        # ============ training one epoch of DINO ... ============
        train_stats = train_one_epoch(student, teacher, teacher_without_ddp, dino_loss,
            data_loader, optimizer, lr_schedule, wd_schedule, momentum_schedule,
            epoch, fp16_scaler, args)

        # ============ writing logs ... ============
        save_dict = {
            'student': student.state_dict(),
            'teacher': teacher.state_dict(),
            'optimizer': optimizer.state_dict(),
            'epoch': epoch + 1,
            'args': args,
            'dino_loss': dino_loss.state_dict(),
        }
        if fp16_scaler is not None:
            save_dict['fp16_scaler'] = fp16_scaler.state_dict()
        utils.save_on_master(save_dict, os.path.join(save_dir, 'checkpoint.pth'))
        if args.saveckp_freq and epoch % args.saveckp_freq == 0:
            utils.save_on_master(save_dict, os.path.join(save_dir, f'checkpoint{str(epoch)}.pth'))
        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                     'epoch': epoch}
        if utils.is_main_process():
            with (Path(save_dir) / "log.txt").open("a") as f:
                f.write(json.dumps(log_stats) + "\n")
    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))


def train_one_epoch(student, teacher, teacher_without_ddp, dino_loss, data_loader,
                    optimizer, lr_schedule, wd_schedule, momentum_schedule,epoch,
                    fp16_scaler, args):
    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Epoch: [{}/{}]'.format(epoch, args.epochs)
    for it, (images, _) in enumerate(metric_logger.log_every(data_loader, 10, header)):

        # update weight decay and learning rate according to their schedule
        it = len(data_loader) * epoch + it  # global training iteration
        for i, param_group in enumerate(optimizer.param_groups):
            param_group["lr"] = lr_schedule[it]
            if i == 0:  # only the first group is regularized
                param_group["weight_decay"] = wd_schedule[it]

        # move images to gpu
        images = [im.cuda(non_blocking=True) for im in images]
        # teacher and student forward passes + compute dino loss
        with torch.cuda.amp.autocast(fp16_scaler is not None):
            teacher_output = teacher(images[:2])  # only the 2 global views pass through the teacher
            student_output = student(images)
            loss = dino_loss(student_output, teacher_output, epoch)

        if not math.isfinite(loss.item()):
            print("Loss is {}, stopping training".format(loss.item()), force=True)
            sys.exit(1)

        # student update
        optimizer.zero_grad()
        param_norms = None
        if fp16_scaler is None:
            loss.backward()
            if args.clip_grad:
                param_norms = utils.clip_gradients(student, args.clip_grad)
            utils.cancel_gradients_last_layer(epoch, student,
                                              args.freeze_last_layer)
            optimizer.step()
        else:
            fp16_scaler.scale(loss).backward()
            if args.clip_grad:
                fp16_scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
                param_norms = utils.clip_gradients(student, args.clip_grad)
            utils.cancel_gradients_last_layer(epoch, student,
                                              args.freeze_last_layer)
            fp16_scaler.step(optimizer)
            fp16_scaler.update()

        # EMA update for the teacher
        with torch.no_grad():
            m = momentum_schedule[it]  # momentum parameter
            for param_q, param_k in zip(student.module.parameters(), teacher_without_ddp.parameters()):
                param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)

        # logging
        torch.cuda.synchronize()
        metric_logger.update(loss=loss.item())
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
        metric_logger.update(wd=optimizer.param_groups[0]["weight_decay"])
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}


class DINOLoss(nn.Module):
    def __init__(self, out_dim, ncrops, warmup_teacher_temp, teacher_temp,
                 warmup_teacher_temp_epochs, nepochs, student_temp=0.1,
                 center_momentum=0.9):
        super().__init__()
        self.student_temp = student_temp
        self.center_momentum = center_momentum
        self.ncrops = ncrops
        self.register_buffer("center", torch.zeros(1, out_dim))
        # we apply a warm up for the teacher temperature because
        # a too high temperature makes the training instable at the beginning
        self.teacher_temp_schedule = np.concatenate((
            np.linspace(warmup_teacher_temp,
                        teacher_temp, warmup_teacher_temp_epochs),
            np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp
        ))

    def forward(self, student_output, teacher_output, epoch):
        """
        Cross-entropy between softmax outputs of the teacher and student networks.
        """
        student_out = student_output / self.student_temp
        student_out = student_out.chunk(self.ncrops)

        # teacher centering and sharpening
        temp = self.teacher_temp_schedule[epoch]
        teacher_out = F.softmax((teacher_output - self.center) / temp, dim=-1)
        teacher_out = teacher_out.detach().chunk(2)

        total_loss = 0
        n_loss_terms = 0
        for iq, q in enumerate(teacher_out):
            for v in range(len(student_out)):
                if v == iq:
                    # we skip cases where student and teacher operate on the same view
                    continue
                loss = torch.sum(-q * F.log_softmax(student_out[v], dim=-1), dim=-1)
                total_loss += loss.mean()
                n_loss_terms += 1
        total_loss /= n_loss_terms
        self.update_center(teacher_output)
        return total_loss

    @torch.no_grad()
    def update_center(self, teacher_output):
        """
        Update center used for teacher output.
        """
        batch_center = torch.sum(teacher_output, dim=0, keepdim=True)
        dist.all_reduce(batch_center)
        batch_center = batch_center / (len(teacher_output) * dist.get_world_size())

        # ema update
        self.center = self.center * self.center_momentum + batch_center * (1 - self.center_momentum)


class DataAugmentationDINO(object):
    def __init__(self,images_are_RGB, global_crops_scale, local_crops_scale, local_crops_number, mean_for_selected_channel,std_for_selected_channel):
        
        if not images_are_RGB:
            flip_gamma_brightness = transforms.Compose([
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.RandomVerticalFlip(p=0.5),
                #utils.AdjustGamma_custom(0.8),
                utils.AdjustBrightness(0.8),
            ])
            normalize = transforms.Compose([
                # utils.normalize_0_to_1(),
                transforms.Normalize(mean_for_selected_channel, std_for_selected_channel),
            ])

            # first global crop
            self.global_transfo1 = transforms.Compose([
                transforms.RandomResizedCrop(224, scale=global_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC),
                flip_gamma_brightness,
                utils.GaussianBlur_forGreyscaleMultiChan(1.0),
                normalize,
            ])
            # second global crop
            self.global_transfo2 = transforms.Compose([
                transforms.RandomResizedCrop(224, scale=global_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC),
                flip_gamma_brightness,
                utils.GaussianBlur_forGreyscaleMultiChan(0.1),
                utils.Solarization_forGreyscaleMultiChan(0.2),
                normalize,
            ])
            # transformation for the local small crops
            self.local_crops_number = local_crops_number
            self.local_transfo = transforms.Compose([
                transforms.RandomResizedCrop(96, scale=local_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC),
                flip_gamma_brightness,
                utils.GaussianBlur_forGreyscaleMultiChan(0.5),
                normalize,
            ])

        #images are RGB
        else:
            flip_and_color_jitter = transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomApply(
                [transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)],
                p=0.8
            ),
            transforms.RandomGrayscale(p=0.2),
            ])
            normalize = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize(mean_for_selected_channel, std_for_selected_channel),
            ])

            # first global crop
            self.global_transfo1 = transforms.Compose([
                transforms.RandomResizedCrop(224, scale=global_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC),
                flip_and_color_jitter,
                utils.GaussianBlur(1.0),
                normalize,
            ])
            # second global crop
            self.global_transfo2 = transforms.Compose([
                transforms.RandomResizedCrop(224, scale=global_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC),
                flip_and_color_jitter,
                utils.GaussianBlur(0.1),
                utils.Solarization(0.2),
                normalize,
            ])
            # transformation for the local small crops
            self.local_crops_number = local_crops_number
            self.local_transfo = transforms.Compose([
                transforms.RandomResizedCrop(96, scale=local_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC),
                flip_and_color_jitter,
                utils.GaussianBlur(p=0.5),
                normalize,
            ])

    def __call__(self, image):
        crops = []
        crops.append(self.global_transfo1(image))
        crops.append(self.global_transfo2(image))
        for _ in range(self.local_crops_number):
            crops.append(self.local_transfo(image))
        for crop in crops:
            if torch.isnan(crop).any():
                print("Nan found in the crop")
        return crops

if __name__ == '__main__':
    parser = argparse.ArgumentParser('DINO', parents=[get_args_parser()])
    args = parser.parse_args()
    
    args.selected_channels = list(map(int, args.selected_channels[0].split(',')))
    args.channel_dict = dict(zip(args.selected_channels, args.channel_dict.split(',')))

    if args.parse_params:
        t_args = argparse.Namespace()
        t_args.__dict__.update(ast.literal_eval(args.parse_params))
        args = parser.parse_args(namespace=t_args)
    save_dir = f"{args.output_dir}/{args.name_of_run}/scDINO_ViTs/{args.full_ViT_name}"
    Path(save_dir).mkdir(parents=True, exist_ok=True)
    train_dino(args, save_dir)

3.test_cases/4.DDP/pyscripts/main_dino.py
aws-samples/awsome-distributed-training
import torch
from torchvision import datasets
import utils
from tifffile import imread
import numpy
import torchvision.transforms as transforms
import json
import yaml
import os

def save_config_file(config, save_dir):
    os.makedirs(save_dir, exist_ok=True)
    with open(f"{save_dir}/run_config_dump.json", "w") as f:
        json.dump(config, f)
    with open(f"{save_dir}/run_config_dump.yaml", "w") as f:
        yaml.dump(config, f)


class ReturnIndexDataset(datasets.ImageFolder):
    def __getitem__(self, idx):
        path, target = self.samples[idx]
        image= imread(path)
        image=image.astype(float)
        if center_crop != 0:
            image = torch.from_numpy(image).permute(2, 0, 1)
            transform = transforms.CenterCrop(center_crop)
            tensor = transform(image)
            image = image.permute(1, 2, 0)
            image = image.numpy()
        image = utils.normalize_numpy_0_to_1(image)
        if utils.check_nan(image):
            print("nan in image: ", path)
            return None
        else:
            tensor = torch.from_numpy(image).permute(2, 0, 1)
            if torch.isnan(tensor).any():
                print("nan in tensor: ", path)
                return None
            else:
                return tensor, idx

def collate_fn(batch):
    batch = list(filter(lambda x: x is not None, batch))
    return torch.utils.data.dataloader.default_collate(batch)

config_file = '../scDINO_full_pipeline.yaml'
with open(config_file, "r") as f:
    config = yaml.load(f,Loader=yaml.FullLoader)

data_dir = config['meta']['dataset_dir'] 
fraction_for_mean_std = config['train_scDINO']['fraction_for_mean_std_calc']
center_crop = config['meta']['center_crop']

name_of_run = config['meta']['name_of_run']
sk_save_dir = config['meta']['output_dir']
save_dir_downstream_run = sk_save_dir+"/"+name_of_run

output_file = f"{save_dir_downstream_run}/mean_and_std_of_dataset.txt"

save_config_file(config, save_dir_downstream_run)

print("loading dataset...")
dataset_total = ReturnIndexDataset(data_dir)
print("length of dataset: ", len(dataset_total))

#SAMPLER SECTION
validation_split = 1-float(fraction_for_mean_std)
shuffle_dataset = True
random_seed= 42
dataset_size = len(dataset_total)
indices = list(range(dataset_size))
split = int(numpy.floor(validation_split * dataset_size))
if shuffle_dataset :
    numpy.random.seed(random_seed)
    numpy.random.shuffle(indices)
train_indices, val_indices = indices[split:], indices[:split]

print(len(train_indices), len(val_indices))

print(f"Train dataset consists of {len(train_indices)} images.")

train_sampler = torch.utils.data.SubsetRandomSampler(train_indices)

image_data_loader = torch.utils.data.DataLoader(
    dataset_total,
    sampler=train_sampler,
    batch_size=int(len(train_indices)/10),
    num_workers=0, 
    collate_fn=collate_fn, 
    drop_last=True)

print("Successfully loaded data.")

def batch_mean_and_sd(loader):
    cnt = 0
    picture, _ = next(iter(image_data_loader))
    b, c, h, w = picture.shape
    fst_moment = torch.empty(c)
    snd_moment = torch.empty(c)
    
    for images, _ in loader:
        b, c, h, w = images.shape
        nb_pixels = b * h * w
        sum_ = torch.sum(images, dim=[0, 2, 3])
        sum_of_square = torch.sum(images ** 2,
                                  dim=[0, 2, 3])
        fst_moment = (cnt * fst_moment + sum_) / (cnt + nb_pixels)
        snd_moment = (cnt * snd_moment + sum_of_square) / (cnt + nb_pixels)
        cnt += nb_pixels

    mean, std = fst_moment, torch.sqrt(snd_moment - fst_moment ** 2)
    return mean,std
  
mean, std = batch_mean_and_sd(image_data_loader)
print("mean and std: \n", mean, std)


with open(output_file, 'w') as f:
    json.dump({"mean": mean.tolist(), "std": std.tolist()}, f)

3.test_cases/4.DDP/pyscripts/mean_std_dataset.py
aws-samples/awsome-distributed-training
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from sklearn import preprocessing
import umap
import os
import seaborn as sns
import yaml
import glob

def get_channel_name_combi(channel_combi_num, channel_dict):
    name_of_channel_combi = ""
    for channel_number in iter(str(channel_combi_num)):
        name_of_channel_combi = "_".join([name_of_channel_combi, channel_dict[int(channel_number)]])
    return name_of_channel_combi

def get_channel_number_combi(channel_names, channel_dict):
    channel_combi = ""
    for channel_name in channel_names.split('_'):
        for key, value in channel_dict.items():
            if value == channel_name:
                channel_combi = "".join([channel_combi, str(key)])
    return channel_combi

def get_channel_name_combi_list(selected_channels, channel_dict):
    channel_names = []
    for channel_combi in selected_channels:
        channel_names.append(get_channel_name_combi(channel_combi,channel_dict))
    return channel_names

config_file = '../scDINO_full_pipeline.yaml'
with open(config_file, "r") as f:
    config = yaml.load(f,Loader=yaml.FullLoader)

name_of_run = config['meta']['name_of_run']
sk_save_dir = config['meta']['output_dir']
save_dir_downstream_run = sk_save_dir+"/"+name_of_run


#channel_aTub_model_checkpoint2_features.csv
features_path = f"{save_dir_downstream_run}/CLS_features/"

features_file = glob.glob(features_path+'*_features.csv')[0]

labels_file = f"{save_dir_downstream_run}/CLS_features/class_labels.csv"


#load data
features = np.genfromtxt(features_file, delimiter = ',')
class_labels_pd = pd.read_csv(labels_file, header=None)
class_labels = class_labels_pd[0].tolist()

#create directories to save plots
#if snakemake.params['scDINO_full_pipeline']:
#    save_dir= f"{snakemake.wildcards['save_dir_downstream_run']}/{snakemake.wildcards['ViT_name']}_channel{snakemake.wildcards['channel_names']}_analyses/embedding_plots/"
#    file_name= f"epoch{snakemake.wildcards['epoch_num']}_"
#    umap_params = snakemake.config['downstream_analyses']['umap_eval']
#else:

dino_vit_name = config['train_scDINO']['dino_vit_name']

selected_channels = [config['meta']['selected_channel_combination_per_run']]
selected_channels = list(eval(selected_channels[0]))

channel_dict = config['meta']['channel_dict']
channel_dict = dict(zip(selected_channels, channel_dict.split(',')))

channel_names=get_channel_name_combi_list(selected_channels, channel_dict)

save_dir = f"{save_dir_downstream_run}/{dino_vit_name}_channel{channel_names[0]}_analyses/embedding_plots/"
file_name = f"{channel_names[0]}_{dino_vit_name}"

#UMAP
def fit_umap(data, n_neighbors, min_dist, metric, spread, epochs):
    umap_model = umap.UMAP(n_neighbors=n_neighbors, min_dist=min_dist, metric=metric, spread=spread, n_epochs=epochs, random_state=42)
    umap_embedding = umap_model.fit_transform(data)
    return umap_embedding

n_neighbors = config['downstream_analyses']['umap_eval']['n_neighbors']
min_dist = config['downstream_analyses']['umap_eval']['min_dist']
metric = config['downstream_analyses']['umap_eval']['metric']
spread = config['downstream_analyses']['umap_eval']['spread']
epochs = config['downstream_analyses']['umap_eval']['epochs']

umap_embedding = fit_umap(features, n_neighbors=15, min_dist=0.1, metric=metric, spread=spread, epochs=epochs)

custom_palette = sns.color_palette("hls", len(set(class_labels)))

def make_plot(embedding, labels, save_dir, file_name=file_name,name="Emb type", description="details"):
    sns_plot = sns.scatterplot(x=embedding[:,0], y=embedding[:,1], hue=labels, s=14, palette=custom_palette, linewidth=0, alpha=0.9)
    plt.suptitle(f"{name}_{file_name}", fontsize=8)
    sns_plot.tick_params(labelbottom=False)
    sns_plot.tick_params(labelleft=False)
    sns_plot.tick_params(bottom=False)
    sns_plot.tick_params(left=False)
    sns_plot.set_title("CLS Token embedding of "+str(len(labels))+" cells with a dimensionality of "+str(features.shape[1])+" \n"+description, fontsize=6)
    sns.move_legend(sns_plot, "lower left", title='Classes', prop={'size': 5}, title_fontsize=6, markerscale=0.5)
    sns.set(rc={"figure.figsize":(14, 10)})
    sns.despine(bottom = True, left = True)
    sns_plot.figure.savefig(f"{save_dir}{file_name}{name}.png", dpi=325)
    sns_plot.figure.savefig(f"{save_dir}pdf_format/{file_name}{name}.pdf")
    plt.close()

os.makedirs(f"{save_dir}pdf_format", exist_ok=True)

make_plot(umap_embedding, class_labels, save_dir=save_dir, file_name=file_name, name="umap",description=f"n_neighbors:{n_neighbors}, min_dist={min_dist}, metric={metric}, spread={spread}, epochs={epochs}")


########################### Additional plots from https://topometry.readthedocs.io/en/latest/ ###########################
if config['downstream_analyses']['umap_eval']['topometry_plots']:
    
    import topo as tp

    os.makedirs(f"{save_dir}/topometry_plots", exist_ok=True)
    os.makedirs(f"{save_dir}/topometry_plots/pdf_format", exist_ok=True)

    save_dir_topo = f"{save_dir}/topometry_plots/"
    # Learn topological metrics and basis from data. The default is to use diffusion harmonics.
    tg = tp.TopOGraph()

    print('running all combinations')
    tg.run_layouts(features, n_components=2,
                        bases=['diffusion', 'fuzzy'],
                        graphs=['diff', 'fuzzy'],
                        layouts=['tSNE', 'MAP', 'MDE', 'PaCMAP', 'TriMAP', 'NCVis'])

    make_plot(tg.db_diff_MAP, class_labels, name="db_diff_MAP", save_dir=save_dir_topo)
    make_plot(tg.db_fuzzy_MAP, class_labels, name="db_fuzzy_MAP", save_dir=save_dir_topo)
    make_plot(tg.db_diff_MDE, class_labels, name="db_diff_MDE", save_dir=save_dir_topo)
    make_plot(tg.db_fuzzy_MDE, class_labels, name="db_fuzzy_MDE", save_dir=save_dir_topo)
    make_plot(tg.db_PaCMAP, class_labels, name="db_PaCMAP", save_dir=save_dir_topo)
    make_plot(tg.db_TriMAP, class_labels, name="db_TriMAP", save_dir=save_dir_topo)
    make_plot(tg.db_tSNE, class_labels, name="db_tSNE", save_dir=save_dir_topo)
    make_plot(tg.fb_diff_MAP, class_labels, name="fb_diff_MAP", save_dir=save_dir_topo)
    make_plot(tg.fb_fuzzy_MAP, class_labels, name="fb_fuzzy_MAP", save_dir=save_dir_topo)
    make_plot(tg.fb_diff_MDE, class_labels, name="fb_diff_MDE", save_dir=save_dir_topo)
    make_plot(tg.fb_fuzzy_MDE, class_labels, name="fb_fuzzy_MDE", save_dir=save_dir_topo)
    make_plot(tg.fb_PaCMAP, class_labels, name="fb_PaCMAP", save_dir=save_dir_topo)
    make_plot(tg.fb_TriMAP, class_labels, name="fb_TriMAP", save_dir=save_dir_topo)
    make_plot(tg.fb_tSNE, class_labels, name="fb_tSNE", save_dir=save_dir_topo)

3.test_cases/4.DDP/pyscripts/plot_in_2D.py
aws-samples/awsome-distributed-training
import os
import re
import ast
import yaml

#### FUNCTIONS ####
def get_channel_name_combi(channel_combi_num, channel_dict):
    name_of_channel_combi = ""
    for channel_number in iter(str(channel_combi_num)):
        name_of_channel_combi = "_".join([name_of_channel_combi, channel_dict[int(channel_number)]])
    return name_of_channel_combi

def get_channel_number_combi(channel_names, channel_dict):
    channel_combi = ""
    for channel_name in channel_names.split('_'):
        for key, value in channel_dict.items():
            if value == channel_name:
                channel_combi = "".join([channel_combi, str(key)])
    return channel_combi

def get_channel_name_combi_list(selected_channels, channel_dict):
    channel_names = []
    for channel_combi in selected_channels:
        channel_names.append(get_channel_name_combi(channel_combi,channel_dict))
    return channel_names

def save_config_file(config, save_dir):
    os.makedirs(save_dir, exist_ok=True)
    with open(f"{save_dir}/run_config_dump.json", "w") as f:
        json.dump(config, f)
    with open(f"{save_dir}/run_config_dump.yaml", "w") as f:
        yaml.dump(config, f)

def load_norm_per_channel(filepath_mean_and_std_of_dataset):
    with open(filepath_mean_and_std_of_dataset) as f:
        norm_per_channel_json = json.load(f)
        norm_per_channel = str([tuple(norm_per_channel_json['mean']), tuple(norm_per_channel_json['std'])])
        return norm_per_channel 


#### PARSING ####
name_of_run = config['meta']['name_of_run']
sk_save_dir = config['meta']['output_dir']
ViT_name = config['train_scDINO']['dino_vit_name']
epochs= int(config['train_scDINO']['epochs']+1)
save_dir_downstream_run = sk_save_dir+"/"+name_of_run
selected_channels = config['meta']['selected_channel_combination_per_run']
channel_dict = config['meta']['channel_dict']
saveckp_freq = int(config['train_scDINO']['saveckp_freq'])
epoch_nums = [epoch_num*saveckp_freq for epoch_num in range(0,(int(epochs/saveckp_freq)))]
if (epochs-1) not in epoch_nums: 
    epoch_nums.append(epochs-1)
print('Epochs for downstream analyses:', epoch_nums)
save_config_file(config, save_dir_downstream_run)

#scDINO-ViTs_path_for_extract_labels
scDINO_ViTs_for_path_extract_labels = f"{save_dir_downstream_run}/scDINO_ViTs/{ViT_name}_channel{get_channel_name_combi_list(selected_channels, channel_dict)[0]}/checkpoint0.pth"

3.test_cases/4.DDP/pyscripts/setup.py
aws-samples/awsome-distributed-training
# Copyright (c) Facebook, Inc. and its affiliates.
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Misc functions.

Mostly copy-paste from torchvision references or other public repos like DETR:
https://github.com/facebookresearch/detr/blob/master/util/misc.py
"""
from cmath import nan
from email import utils
import os
import sys
import time
import math
import random
import datetime
import subprocess
from collections import defaultdict, deque

import numpy as np
import torch
from torch import nn
import torch.distributed as dist
from PIL import ImageFilter, ImageOps
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
import random
import warnings
from datetime import timedelta

warnings.filterwarnings('ignore')


def collate_fn(batch):
    batch = list(filter(lambda x: x is not None, batch))
    return torch.utils.data.dataloader.default_collate(batch)


def check_nan(x):
    if np.isnan(x).any():
        return True
    else:
        return False


def check_zero(x):
    x_0 = np.where(x == 0)
    if len(x_0[0]) > 0:
        print("x_0: {}".format(x_0))
        return True
    

def normalize_tensor_per_channel(x):
    x_min = x.min(axis=(1,2), keepdims=True)
    x_max = x.max(axis=(1,2), keepdims=True)
    diff_min_max = x_max - x_min
    if check_nan(diff_min_max):
        print("diff_min_max is nan")
    if check_nan(x-x_min):
        print("x-x_min is nan:")
    if check_nan(x):
        print("x contains nan before normalization")
    if check_zero(diff_min_max):
        print("diff_min_max has zero")
        print("x_max",x_max)
        print("x_min",x_min)
        print("diff_min_max",diff_min_max)
        print("x",x.shape)
    x = (x - x_min)/(x_max-x_min)
    if check_nan(x):
        print("x contains nan after normalization")
    return x


class normalize_tensor_0_to_1(object):
    def __call__(self, tensor):
        image = tensor.numpy()
        image = normalize_tensor_per_channel(image)
        tensor = torch.from_numpy(image).float()
        return tensor


def normalize_numpy_0_to_1(x):
    x_min = x.min(axis=(0,1), keepdims=True)
    x_max = x.max(axis=(0,1), keepdims=True)
    diff_min_max = x_max - x_min
    if check_nan(diff_min_max):
        print("diff_min_max is nan")
    if check_nan(x-x_min):
        print("x-x_min is nan:")
    if check_nan(x):
        print("x contains nan before normalization")
    if check_zero(diff_min_max):
        print("diff_min_max has zero")
        print("x_max",x_max)
        print("x_min",x_min)
        print("diff_min_max",diff_min_max)
        print("x",x.shape)
    x = (x - x_min)/(x_max-x_min)
    if check_nan(x):
        print("x contains nan after normalization")
    return x


class GaussianBlur(object):
    """
    Apply Gaussian Blur to the PIL image.
    """
    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.):
        self.prob = p
        self.radius_min = radius_min
        self.radius_max = radius_max

    def __call__(self, img):
        do_it = random.random() <= self.prob
        if not do_it:
            return img

        return img.filter(
            ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)
            )
        )

class GaussianBlur_forGreyscaleMultiChan(object):
    """
    Apply Gaussian Blur to a Tensor.
    """
    def __init__(self, p=0.5, sigma_min=0.1, sigma_max=2.):
        self.prob = p
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max

    def __call__(self, tensor):
        do_it = random.random() <= self.prob
        if not do_it:
            return tensor

        else:
            blur = transforms.GaussianBlur(kernel_size=(5, 5), sigma=(0.1, 1.5))
            tensor = blur(tensor)
            return tensor


class Solarization(object):
    """
    Apply Solarization to the PIL image.
    """
    def __init__(self, p):
        self.p = p

    def __call__(self, img):
        if random.random() < self.p:
            return ImageOps.solarize(img)
        else:
            return img

class Solarization_forGreyscaleMultiChan(object):
    """
    Apply Solarization to the Tensor.
    """
    def __init__(self, p, threshold=0.5):
        self.p = p
        self.threshold = threshold

    def __call__(self, tensor):
        if random.random() < self.p:
            tensor[tensor > self.threshold] = 1 - tensor[tensor > self.threshold]
            return tensor
        else:
            return tensor


class AdjustBrightness(object): 
    def __init__(self, p):
        self.p = p
    def __call__(self, x):
        if random.random() < self.p:
            for channel in range(x.shape[0]):
                x[channel] = TF.adjust_brightness(x[channel], random.uniform(0.5,1.25))
            return x
        else:
            return x

class AdjustGamma_custom(object):
    def __init__(self, p, eps=1e-8):
        self.p = p
        self.eps = eps
    def __call__(self, x):
        if random.random() < self.p:
            random_gamma= random.uniform(0.75,3)
            for channel in range(x.shape[0]):
                x[channel] = gamma_adjust_wo_nan(x[channel],random_gamma, eps=self.eps)
            return x
        else:
            return x


def fetch_foldername_of_img_location_donor(dataset, index):
    img_path = dataset.imgs[index][0]
    img_path = img_path.split("/")
    img_folder_name = img_path[-4]
    return img_folder_name

def fetch_foldername_of_img_location(dataset, index, folder_depth=0):
    return dataset.imgs[index][0].split("/")[(-folder_depth-2)]

def list_folders_in_directory(directory):
    return [f.name for f in os.scandir(directory) if f.is_dir() ]

def fetch_all_folder_names_of_folder_depth(dataset, index=1, folder_depth=0):
    path = dataset.imgs[index][0]
    path = "/".join(path.split("/")[:(-folder_depth-2)])
    class_names = list_folders_in_directory(path)
    return class_names


#adapted from https://discuss.pytorch.org/t/problem-training-gamma-correction/60800/2
# There are two tricks here. The first is to handle negative inputs by odd symmetry 6; f(-x) is computed as -f(x). The second is to handle zero or close-to-zero inputs by adding a small ‘epsilon’ value to ensure numerical stability.
# I tested this layer and was able to get it to train whereas without these tricks the self.gamma parameter would end up NaN.
def gamma_adjust_wo_nan(tensor,gamma, eps=1e-8, a=1):
    channel = tensor.numpy()
    gamma_adj_channel = ((np.abs(channel)  + eps)**gamma)*a*np.sign(channel)
    tensor_adj = torch.from_numpy(gamma_adj_channel).float()
    return tensor_adj

def get_pretrained_weights_in_chans(pretrained_weights, checkpoint_key="teacher"):
     if os.path.isfile(pretrained_weights):
        state_dict = torch.load(pretrained_weights, map_location="cpu")
        if checkpoint_key is not None and checkpoint_key in state_dict:
            print(f"Take key {checkpoint_key} in provided checkpoint dict")
            state_dict_check = state_dict[checkpoint_key]
            num_in_chans = state_dict_check["backbone.patch_embed.proj.weight"].shape[1]
        else:
            try:
                num_in_chans = state_dict["patch_embed.proj.weight"].shape[1]
            # print the error of the exception
            except Exception as e:
                print(e)
                print("in_chans not found in checkpoint dict.", state_dict.keys())
                print("Please provide a checkpoint with in_chans in the state_dict.")
                print("Exiting...")
                sys.exit(1)
        return num_in_chans

def load_pretrained_weights(model, pretrained_weights, checkpoint_key, model_name, patch_size):
    if os.path.isfile(pretrained_weights):
        state_dict = torch.load(pretrained_weights, map_location="cpu")
        if checkpoint_key is not None and checkpoint_key in state_dict:
            print(f"Take key {checkpoint_key} in provided checkpoint dict")
            state_dict = state_dict[checkpoint_key]
        # remove `module.` prefix
        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
        # remove `backbone.` prefix induced by multicrop wrapper
        state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
        msg = model.load_state_dict(state_dict, strict=False)
        print('Pretrained weights found at {} and loaded with msg: {}'.format(pretrained_weights, msg))
    else:
        print("Please use the `--pretrained_weights` argument to indicate the path of the checkpoint to evaluate.")
        url = None
        if model_name == "vit_small" and patch_size == 16:
            url = "dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth"
        elif model_name == "vit_small" and patch_size == 8:
            url = "dino_deitsmall8_pretrain/dino_deitsmall8_pretrain.pth"
        elif model_name == "vit_base" and patch_size == 16:
            url = "dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth"
        elif model_name == "vit_base" and patch_size == 8:
            url = "dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth"
        elif model_name == "xcit_small_12_p16":
            url = "dino_xcit_small_12_p16_pretrain/dino_xcit_small_12_p16_pretrain.pth"
        elif model_name == "xcit_small_12_p8":
            url = "dino_xcit_small_12_p8_pretrain/dino_xcit_small_12_p8_pretrain.pth"
        elif model_name == "xcit_medium_24_p16":
            url = "dino_xcit_medium_24_p16_pretrain/dino_xcit_medium_24_p16_pretrain.pth"
        elif model_name == "xcit_medium_24_p8":
            url = "dino_xcit_medium_24_p8_pretrain/dino_xcit_medium_24_p8_pretrain.pth"
        elif model_name == "resnet50":
            url = "dino_resnet50_pretrain/dino_resnet50_pretrain.pth"
        else:
            print("There is no reference weights available for this model => We use random weights.")


def load_pretrained_linear_weights(linear_classifier, model_name, patch_size):
    url = None
    if model_name == "vit_small" and patch_size == 16:
        url = "dino_deitsmall16_pretrain/dino_deitsmall16_linearweights.pth"
    elif model_name == "vit_small" and patch_size == 8:
        url = "dino_deitsmall8_pretrain/dino_deitsmall8_linearweights.pth"
    elif model_name == "vit_base" and patch_size == 16:
        url = "dino_vitbase16_pretrain/dino_vitbase16_linearweights.pth"
    elif model_name == "vit_base" and patch_size == 8:
        url = "dino_vitbase8_pretrain/dino_vitbase8_linearweights.pth"
    elif model_name == "resnet50":
        url = "dino_resnet50_pretrain/dino_resnet50_linearweights.pth"
    if url is not None:
        print("We load the reference pretrained linear weights.")
        state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/" + url)["state_dict"]
        linear_classifier.load_state_dict(state_dict, strict=True)
    else:
        print("We use random linear weights.")


def clip_gradients(model, clip):
    norms = []
    for name, p in model.named_parameters():
        if p.grad is not None:
            param_norm = p.grad.data.norm(2)
            norms.append(param_norm.item())
            clip_coef = clip / (param_norm + 1e-6)
            if clip_coef < 1:
                p.grad.data.mul_(clip_coef)
    return norms


def cancel_gradients_last_layer(epoch, model, freeze_last_layer):
    if epoch >= freeze_last_layer:
        return
    for n, p in model.named_parameters():
        if "last_layer" in n:
            p.grad = None


def restart_from_checkpoint(ckp_path, run_variables=None, **kwargs):
    """
    Re-start from checkpoint
    """
    if not os.path.isfile(ckp_path):
        return
    print("Found checkpoint at {}".format(ckp_path))

    # open checkpoint file
    checkpoint = torch.load(ckp_path, map_location="cpu")

    # key is what to look for in the checkpoint file
    # value is the object to load
    # example: {'state_dict': model}
    for key, value in kwargs.items():
        if key in checkpoint and value is not None:
            try:
                msg = value.load_state_dict(checkpoint[key], strict=False)
                print("=> loaded '{}' from checkpoint '{}' with msg {}".format(key, ckp_path, msg))
            except TypeError:
                try:
                    msg = value.load_state_dict(checkpoint[key])
                    print("=> loaded '{}' from checkpoint: '{}'".format(key, ckp_path))
                except ValueError:
                    print("=> failed to load '{}' from checkpoint: '{}'".format(key, ckp_path))
        else:
            print("=> key '{}' not found in checkpoint: '{}'".format(key, ckp_path))

    # re load variable important for the run
    if run_variables is not None:
        for var_name in run_variables:
            if var_name in checkpoint:
                run_variables[var_name] = checkpoint[var_name]


def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0, start_warmup_value=0):
    warmup_schedule = np.array([])
    warmup_iters = warmup_epochs * niter_per_ep
    if warmup_epochs > 0:
        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)

    iters = np.arange(epochs * niter_per_ep - warmup_iters)
    schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))

    schedule = np.concatenate((warmup_schedule, schedule))
    assert len(schedule) == epochs * niter_per_ep
    return schedule


def bool_flag(s):
    """
    Parse boolean arguments from the command line.
    """
    FALSY_STRINGS = {"off", "false", "0"}
    TRUTHY_STRINGS = {"on", "true", "1"}
    if s.lower() in FALSY_STRINGS:
        return False
    elif s.lower() in TRUTHY_STRINGS:
        return True
    else:
        raise argparse.ArgumentTypeError("invalid value for a boolean flag")


def fix_random_seeds(seed=31):
    """
    Fix random seeds.
    """
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)


class SmoothedValue(object):
    """Track a series of values and provide access to smoothed values over a
    window or the global series average.
    """

    def __init__(self, window_size=20, fmt=None):
        if fmt is None:
            fmt = "{median:.6f} ({global_avg:.6f})"
        self.deque = deque(maxlen=window_size)
        self.total = 0.0
        self.count = 0
        self.fmt = fmt

    def update(self, value, n=1):
        self.deque.append(value)
        self.count += n
        self.total += value * n

    def synchronize_between_processes(self):
        """
        Warning: does not synchronize the deque!
        """
        if not is_dist_avail_and_initialized():
            return
        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
        dist.barrier()
        dist.all_reduce(t)
        t = t.tolist()
        self.count = int(t[0])
        self.total = t[1]

    @property
    def median(self):
        d = torch.tensor(list(self.deque))
        return d.median().item()

    @property
    def avg(self):
        d = torch.tensor(list(self.deque), dtype=torch.float32)
        return d.mean().item()

    @property
    def global_avg(self):
        return self.total / self.count

    @property
    def max(self):
        return max(self.deque)

    @property
    def value(self):
        return self.deque[-1]

    def __str__(self):
        return self.fmt.format(
            median=self.median,
            avg=self.avg,
            global_avg=self.global_avg,
            max=self.max,
            value=self.value)


def reduce_dict(input_dict, average=True):
    """
    Args:
        input_dict (dict): all the values will be reduced
        average (bool): whether to do average or sum
    Reduce the values in the dictionary from all processes so that all processes
    have the averaged results. Returns a dict with the same fields as
    input_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return input_dict
    with torch.no_grad():
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        for k in sorted(input_dict.keys()):
            names.append(k)
            values.append(input_dict[k])
        values = torch.stack(values, dim=0)
        dist.all_reduce(values)
        if average:
            values /= world_size
        reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict


class MetricLogger(object):
    def __init__(self, delimiter="\t"):
        self.meters = defaultdict(SmoothedValue)
        self.delimiter = delimiter

    def update(self, **kwargs):
        for k, v in kwargs.items():
            if isinstance(v, torch.Tensor):
                v = v.item()
            assert isinstance(v, (float, int))
            self.meters[k].update(v)

    def __getattr__(self, attr):
        if attr in self.meters:
            return self.meters[attr]
        if attr in self.__dict__:
            return self.__dict__[attr]
        raise AttributeError("'{}' object has no attribute '{}'".format(
            type(self).__name__, attr))

    def __str__(self):
        loss_str = []
        for name, meter in self.meters.items():
            loss_str.append(
                "{}: {}".format(name, str(meter))
            )
        return self.delimiter.join(loss_str)

    def synchronize_between_processes(self):
        for meter in self.meters.values():
            meter.synchronize_between_processes()

    def add_meter(self, name, meter):
        self.meters[name] = meter

    def log_every(self, iterable, print_freq, header=None):
        i = 0
        if not header:
            header = ''
        start_time = time.time()
        end = time.time()
        iter_time = SmoothedValue(fmt='{avg:.6f}')
        data_time = SmoothedValue(fmt='{avg:.6f}')
        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
        if torch.cuda.is_available():
            log_msg = self.delimiter.join([
                header,
                '[{0' + space_fmt + '}/{1}]',
                'eta: {eta}',
                '{meters}',
                'time: {time}',
                'data: {data}',
                'max mem: {memory:.0f}'
            ])
        else:
            log_msg = self.delimiter.join([
                header,
                '[{0' + space_fmt + '}/{1}]',
                'eta: {eta}',
                '{meters}',
                'time: {time}',
                'data: {data}'
            ])
        MB = 1024.0 * 1024.0
        for obj in iterable:
            data_time.update(time.time() - end)
            yield obj
            iter_time.update(time.time() - end)
            if i % print_freq == 0 or i == len(iterable) - 1:
                eta_seconds = iter_time.global_avg * (len(iterable) - i)
                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                if torch.cuda.is_available():
                    print(log_msg.format(
                        i, len(iterable), eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time),
                        memory=torch.cuda.max_memory_allocated() / MB))
                else:
                    print(log_msg.format(
                        i, len(iterable), eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time)))
            i += 1
            end = time.time()
        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
        print('{} Total time: {} ({:.6f} s / it)'.format(
            header, total_time_str, total_time / len(iterable)))


def get_sha():
    cwd = os.path.dirname(os.path.abspath(__file__))

    def _run(command):
        return subprocess.check_output(command, cwd=cwd).decode('ascii').strip()
    sha = 'N/A'
    diff = "clean"
    branch = 'N/A'
    try:
        sha = _run(['git', 'rev-parse', 'HEAD'])
        subprocess.check_output(['git', 'diff'], cwd=cwd)
        diff = _run(['git', 'diff-index', 'HEAD'])
        diff = "has uncommited changes" if diff else "clean"
        branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])
    except Exception:
        pass
    message = f"sha: {sha}, status: {diff}, branch: {branch}"
    return message


def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True


def get_world_size():
    if not is_dist_avail_and_initialized():
        return 1
    return dist.get_world_size()


def get_rank():
    if not is_dist_avail_and_initialized():
        return 0
    return dist.get_rank()


def is_main_process():
    return get_rank() == 0


def save_on_master(*args, **kwargs):
    if is_main_process():
        torch.save(*args, **kwargs)


def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print


def init_distributed_mode(args):
    # launched with torch.distributed.launch
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        args.gpu = int(os.environ['LOCAL_RANK'])
    # launched with submitit on a slurm cluster
    elif 'SLURM_PROCID' in os.environ:
        args.rank = int(os.environ['SLURM_PROCID'])
        args.gpu = args.rank % torch.cuda.device_count()
    # launched naively with `python main_dino.py`
    # we manually add MASTER_ADDR and MASTER_PORT to env variables
    elif torch.cuda.is_available():
        print('Will run the code on one GPU.')
        args.rank, args.gpu, args.world_size = 0, 0, 1
        # os.environ['MASTER_ADDR'] = '127.0.0.1'
        #define another random adress to avoid address already in use error
        os.environ['MASTER_ADDR'] = '127.0.0.{}'.format(random.randint(1, 255))
        # os.environ['MASTER_PORT'] = '29500'
        os.environ['MASTER_PORT'] = str(random.randint(29500, 30000))
        print('MASTER_ADDR: {}'.format(os.environ['MASTER_ADDR']))
    else:
        print('Does not support training without GPU.')
        sys.exit(1)

    dist.init_process_group(
        backend="nccl",
        init_method=args.dist_url,
        world_size=args.world_size,
        rank=args.rank,
        timeout=timedelta(seconds=120)
    )

    torch.cuda.set_device(args.gpu)
    print('| distributed init (rank {}): {}'.format(
        args.rank, args.dist_url), flush=True)
    dist.barrier()
    setup_for_distributed(args.rank == 0)


def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)
    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.reshape(1, -1).expand_as(pred))
    return [correct[:k].reshape(-1).float().sum(0) * 100. / batch_size for k in topk]


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


class LARS(torch.optim.Optimizer):
    """
    Almost copy-paste from https://github.com/facebookresearch/barlowtwins/blob/main/main.py
    """
    def __init__(self, params, lr=0, weight_decay=0, momentum=0.9, eta=0.001,
                 weight_decay_filter=None, lars_adaptation_filter=None):
        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum,
                        eta=eta, weight_decay_filter=weight_decay_filter,
                        lars_adaptation_filter=lars_adaptation_filter)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            for p in g['params']:
                dp = p.grad

                if dp is None:
                    continue

                if p.ndim != 1:
                    dp = dp.add(p, alpha=g['weight_decay'])

                if p.ndim != 1:
                    param_norm = torch.norm(p)
                    update_norm = torch.norm(dp)
                    one = torch.ones_like(param_norm)
                    q = torch.where(param_norm > 0.,
                                    torch.where(update_norm > 0,
                                                (g['eta'] * param_norm / update_norm), one), one)
                    dp = dp.mul(q)

                param_state = self.state[p]
                if 'mu' not in param_state:
                    param_state['mu'] = torch.zeros_like(p)
                mu = param_state['mu']
                mu.mul_(g['momentum']).add_(dp)

                p.add_(mu, alpha=-g['lr'])


class MultiCropWrapper(nn.Module):
    """
    Perform forward pass separately on each resolution input.
    The inputs corresponding to a single resolution are clubbed and single
    forward is run on the same resolution inputs. Hence we do several
    forward passes = number of different resolutions used. We then
    concatenate all the output features and run the head forward on these
    concatenated features.
    """
    def __init__(self, backbone, head):
        super(MultiCropWrapper, self).__init__()
        # disable layers dedicated to ImageNet labels classification
        backbone.fc, backbone.head = nn.Identity(), nn.Identity()
        self.backbone = backbone
        self.head = head

    def forward(self, x):
        # convert to list
        if not isinstance(x, list):
            x = [x]
        idx_crops = torch.cumsum(torch.unique_consecutive(
            torch.tensor([inp.shape[-1] for inp in x]),
            return_counts=True,
        )[1], 0)
        start_idx, output = 0, torch.empty(0).to(x[0].device)
        for end_idx in idx_crops:
            _out = self.backbone(torch.cat(x[start_idx: end_idx]))
            # The output is a tuple with XCiT model. See:
            # https://github.com/facebookresearch/xcit/blob/master/xcit.py#L404-L405
            if isinstance(_out, tuple):
                _out = _out[0]
            # accumulate outputs
            output = torch.cat((output, _out))
            start_idx = end_idx
        # Run the head forward on the concatenated features.
        return self.head(output)


def get_params_groups(model):
    regularized = []
    not_regularized = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        # we do not regularize biases nor Norm parameters
        if name.endswith(".bias") or len(param.shape) == 1:
            not_regularized.append(param)
        else:
            regularized.append(param)
    return [{'params': regularized}, {'params': not_regularized, 'weight_decay': 0.}]


def has_batchnorms(model):
    bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
    for name, module in model.named_modules():
        if isinstance(module, bn_types):
            return True
    return False


class PCA():
    """
    Class to  compute and apply PCA.
    """
    def __init__(self, dim=256, whit=0.5):
        self.dim = dim
        self.whit = whit
        self.mean = None

    def train_pca(self, cov):
        """
        Takes a covariance matrix (np.ndarray) as input.
        """
        d, v = np.linalg.eigh(cov)
        eps = d.max() * 1e-5
        n_0 = (d < eps).sum()
        if n_0 > 0:
            d[d < eps] = eps

        # total energy
        totenergy = d.sum()

        # sort eigenvectors with eigenvalues order
        idx = np.argsort(d)[::-1][:self.dim]
        d = d[idx]
        v = v[:, idx]

        print("keeping %.2f %% of the energy" % (d.sum() / totenergy * 100.0))

        # for the whitening
        d = np.diag(1. / d**self.whit)

        # principal components
        self.dvt = np.dot(d, v.T)

    def apply(self, x):
        # input is from numpy
        if isinstance(x, np.ndarray):
            if self.mean is not None:
                x -= self.mean
            return np.dot(self.dvt, x.T).T

        # input is from torch and is on GPU
        if x.is_cuda:
            if self.mean is not None:
                x -= torch.cuda.FloatTensor(self.mean)
            return torch.mm(torch.cuda.FloatTensor(self.dvt), x.transpose(0, 1)).transpose(0, 1)

        # input if from torch, on CPU
        if self.mean is not None:
            x -= torch.FloatTensor(self.mean)
        return torch.mm(torch.FloatTensor(self.dvt), x.transpose(0, 1)).transpose(0, 1)


def compute_ap(ranks, nres):
    """
    Computes average precision for given ranked indexes.
    Arguments
    ---------
    ranks : zerro-based ranks of positive images
    nres  : number of positive images
    Returns
    -------
    ap    : average precision
    """

    # number of images ranked by the system
    nimgranks = len(ranks)

    # accumulate trapezoids in PR-plot
    ap = 0

    recall_step = 1. / nres

    for j in np.arange(nimgranks):
        rank = ranks[j]

        if rank == 0:
            precision_0 = 1.
        else:
            precision_0 = float(j) / rank

        precision_1 = float(j + 1) / (rank + 1)

        ap += (precision_0 + precision_1) * recall_step / 2.

    return ap


def compute_map(ranks, gnd, kappas=[]):
    """
    Computes the mAP for a given set of returned results.
         Usage:
           map = compute_map (ranks, gnd)
                 computes mean average precsion (map) only
           map, aps, pr, prs = compute_map (ranks, gnd, kappas)
                 computes mean average precision (map), average precision (aps) for each query
                 computes mean precision at kappas (pr), precision at kappas (prs) for each query
         Notes:
         1) ranks starts from 0, ranks.shape = db_size X #queries
         2) The junk results (e.g., the query itself) should be declared in the gnd stuct array
         3) If there are no positive images for some query, that query is excluded from the evaluation
    """

    map = 0.
    nq = len(gnd) # number of queries
    aps = np.zeros(nq)
    pr = np.zeros(len(kappas))
    prs = np.zeros((nq, len(kappas)))
    nempty = 0

    for i in np.arange(nq):
        qgnd = np.array(gnd[i]['ok'])

        # no positive images, skip from the average
        if qgnd.shape[0] == 0:
            aps[i] = float('nan')
            prs[i, :] = float('nan')
            nempty += 1
            continue

        try:
            qgndj = np.array(gnd[i]['junk'])
        except:
            qgndj = np.empty(0)

        # sorted positions of positive and junk images (0 based)
        pos  = np.arange(ranks.shape[0])[np.in1d(ranks[:,i], qgnd)]
        junk = np.arange(ranks.shape[0])[np.in1d(ranks[:,i], qgndj)]

        k = 0;
        ij = 0;
        if len(junk):
            # decrease positions of positives based on the number of
            # junk images appearing before them
            ip = 0
            while (ip < len(pos)):
                while (ij < len(junk) and pos[ip] > junk[ij]):
                    k += 1
                    ij += 1
                pos[ip] = pos[ip] - k
                ip += 1

        # compute ap
        ap = compute_ap(pos, len(qgnd))
        map = map + ap
        aps[i] = ap

        # compute precision @ k
        pos += 1 # get it to 1-based
        for j in np.arange(len(kappas)):
            kq = min(max(pos), kappas[j]); 
            prs[i, j] = (pos <= kq).sum() / kq
        pr = pr + prs[i, :]

    map = map / (nq - nempty)
    pr = pr / (nq - nempty)

    return map, aps, pr, prs


def multi_scale(samples, model):
    v = None
    for s in [1, 1/2**(1/2), 1/2]:  # we use 3 different scales
        if s == 1:
            inp = samples.clone()
        else:
            inp = nn.functional.interpolate(samples, scale_factor=s, mode='bilinear', align_corners=False)
        feats = model(inp).clone()
        if v is None:
            v = feats
        else:
            v += feats
    v /= 3
    v /= v.norm()
    return v

3.test_cases/4.DDP/pyscripts/utils.py
aws-samples/awsome-distributed-training
import argparse
import utils
from tifffile import imread
import numpy as np
import ast
import torch
from torchvision import transforms, datasets
import os
from torchvision import models as torchvision_models
from torch import nn
import re
import sys
import cv2
import random
import colorsys
import skimage.io
from skimage.measure import find_contours
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
import torch.nn as nn
import torchvision
import vision_transformer as vits
import numpy as np
import json


def apply_mask(image, mask, color, alpha=0.5):
    for c in range(3):
        image[:, :, c] = image[:, :, c] * (1 - alpha * mask) + alpha * mask * color[c] * 255
    return image


def random_colors(N, bright=True):
    """
    Generate random colors.
    """
    brightness = 1.0 if bright else 0.7
    hsv = [(i / N, 1, brightness) for i in range(N)]
    colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))
    random.shuffle(colors)
    return colors


def display_instances(image, mask, fname="test", figsize=(5, 5), blur=False, contour=True, alpha=0.5):
    fig = plt.figure(figsize=figsize, frameon=False)
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)
    ax = plt.gca()

    N = 1
    mask = mask[None, :, :]
    # Generate random colors
    colors = random_colors(N)

    # Show area outside image boundaries.
    height, width = image.shape[:2]
    margin = 0
    ax.set_ylim(height + margin, -margin)
    ax.set_xlim(-margin, width + margin)
    ax.axis('off')
    masked_image = image.astype(np.uint32).copy()
    for i in range(N):
        color = colors[i]
        _mask = mask[i]
        if blur:
            _mask = cv2.blur(_mask,(10,10))
        # Mask
        masked_image = apply_mask(masked_image, _mask, color, alpha)
        # Mask Polygon
        # Pad to ensure proper polygons for masks that touch image edges.
        if contour:
            padded_mask = np.zeros((_mask.shape[0] + 2, _mask.shape[1] + 2))
            padded_mask[1:-1, 1:-1] = _mask
            contours = find_contours(padded_mask, 0.5)
            for verts in contours:
                # Subtract the padding and flip (y, x) to (x, y)
                verts = np.fliplr(verts) - 1
                p = Polygon(verts, facecolor="none", edgecolor=color)
                ax.add_patch(p)
    ax.imshow(masked_image.astype(np.uint8), aspect='auto')
    fig.savefig(fname)
    return


if __name__ == '__main__':
    parser = argparse.ArgumentParser('Visualize Self-Attention maps')
    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
    parser.add_argument('--patch_size', default=8, type=int, help='Patch resolution of the model.')
    parser.add_argument('--pretrained_weights', default='', type=str,
        help="Path to pretrained weights to load.")
    parser.add_argument("--checkpoint_key", default="teacher", type=str,
        help='Key to use in the checkpoint (example: "teacher")')
    parser.add_argument('--resize_attention_image',type=str, default=True, help="Wheter to resize image or not")
    parser.add_argument("--image_size", default=(480, 480), type=int, nargs="+", help="Resize image.")
    parser.add_argument("--threshold", type=float, default=None, help="""We visualize masks
        obtained by thresholding the self-attention maps to keep xx% of the mass.""")
    parser.add_argument('--output_dir', default='.', type=str)    
    parser.add_argument('--dataset_dir', default='/path/to/imagenet/', type=str)
    parser.add_argument ('--images_are_RGB',help='If images are RGB, set this to True. If images are grayscale, set this to False.', default=False, type=utils.bool_flag)
    parser.add_argument('--selected_channels', default=[0,1,2], nargs='+', help="""list of channel indexes of the .tiff images which should be used to create the tensors.""")
    parser.add_argument('--channel_dict', default=None, type=str,help="""name of the channels in format as dict channel_number, channel_name.""")
    parser.add_argument('--resize', default=False, help="""if images should be resized""")
    parser.add_argument('--resize_length', default=None, help="""quadratic resize length to resize images""")
    parser.add_argument('--norm_per_channel', default="[(x, x, x, x, x), (x, x, x, x, x)]", type=str, help="""2x tuple of mean and std per channel typically values between 0 and 1""")
    parser.add_argument('--norm_per_channel_file', default=None, help="""path to file with mean and std per channel in json format.""")
    parser.add_argument('--center_crop', default=None, help="""center crop factor to crop images""")
    parser.add_argument('--normalize', default="False", type=str, help="""normalize with mean and std per channel""")
    parser.add_argument('--patch_embedding_mapping', default=None, help="""change the patch embedding weights by inputting a string of the sequence of rearrangement of the model '[0,1,2]' or the string 'average_weights' or None""")
    parser.add_argument('--parse_params',help='Load settings from file in json format. Command line options override values in file.')
    parser.add_argument('--use_weighted_sampler', default=False, type=bool, help='Use weighted sampler for training.')
    parser.add_argument("--num_images", default=None, type=int, help="Number of images to run in total.")
    parser.add_argument("--read_model_arch_dynamically", default=None, type=str, help="Read model architecture from pretrained weights")
    parser.add_argument("--use_mean_patch_embedding", default=False, type=bool, help="Use mean patch embedding instead of first patch embedding")
    parser.add_argument("--use_custom_embedding_map", default=False, type=bool, help="Use custom embedding map")
    parser.add_argument("--custom_embedding_map", default=None, type=dict, help="Custom embedding map")
    parser.add_argument("--name_of_run", default=False, type=bool, help="Name of job")
    parser.add_argument("--num_per_class", default=False, type=bool, help="Use custom embedding map")
    parser.add_argument("--scDINO_full_pipeline", default=False, type=bool, help="Using scDINO full pipeline")
    parser.add_argument('--full_ViT_name', default='full_vit_name', type=str, help='name channel combi ViT')
    parser.add_argument('--seed', default=42, type=int, help='Random seed.')

    #save settings
    args = parser.parse_args()

    args.selected_channels = list(map(int, args.selected_channels[0].split(',')))
    args.channel_dict = dict(zip(args.selected_channels, args.channel_dict.split(',')))

    if args.parse_params:
        t_args = argparse.Namespace()
        t_args.__dict__.update(ast.literal_eval(args.parse_params))
        args = parser.parse_args(namespace=t_args)

    # read mean and std per channel
    if args.norm_per_channel_file:
        with open(args.norm_per_channel_file) as f:
            norm_per_channel_json = json.load(f)
            norm_per_channel = str([tuple(norm_per_channel_json['mean']), tuple(norm_per_channel_json['std'])])
            args.norm_per_channel = norm_per_channel

    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))

#adjust model arch and patch size according to pretrained weights
def adjust_model_architecture(args):
    model_name = args.pretrained_weights.split("/")[-1]
    number = re.findall(r'\d+', model_name)
    if number:
        args.patch_size = int(number[0])
    if "small" in model_name:
        args.arch = "vit_small"
    elif "base" in model_name:
        args.arch = "vit_base"

if args.read_model_arch_dynamically:
    adjust_model_architecture(args)

# ============ building network ... ============
# build model
device = torch.device("cpu")
selected_channels = list(map(int, args.selected_channels))
num_in_chans_pretrained = utils.get_pretrained_weights_in_chans(args.pretrained_weights)
model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0, in_chans=num_in_chans_pretrained)
for p in model.parameters():
    p.requires_grad = False
model.eval()
model.to(device)
print('device:', device)
if os.path.isfile(args.pretrained_weights):
    state_dict = torch.load(args.pretrained_weights, map_location="cpu")
    if args.checkpoint_key is not None and args.checkpoint_key in state_dict:
        print(f"Take key {args.checkpoint_key} in provided checkpoint dict")
        state_dict = state_dict[args.checkpoint_key]
    # remove `module.` prefix
    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
    # remove `backbone.` prefix induced by multicrop wrapper
    state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
    msg = model.load_state_dict(state_dict, strict=False)
    print('Pretrained weights found at {} and loaded with msg: {}'.format(args.pretrained_weights, msg))
else:
    print("Please use the `--pretrained_weights` argument to indicate the path of the checkpoint to evaluate.")

# ============ adjusting ViT model option... ============

def embedding_seq(custom_embedding_map, selected_channels):
        map_dict = ast.literal_eval(custom_embedding_map)
        map_dict = {int(k):int(v) for k,v in map_dict.items()}
        embedding_seq = [map_dict[int(input_channel)] for input_channel in selected_channels]
        return embedding_seq

def build_weight_emb(embedding_seq, model):
    weights = model.patch_embed.proj.weight
    weights = weights[:,embedding_seq,:,:]
    model.patch_embed.proj.weight = nn.Parameter(weights)
    return model

if not args.use_mean_patch_embedding and not args.use_custom_embedding_map:
    if len(selected_channels) != num_in_chans_pretrained:
        print(f"Error: Number of channels in the dataset ({len(selected_channels)}) and pretrained weights ({num_in_chans_pretrained}) are different")
        print(f"Use --use_mean_patch_embedding or --use_custom_embedding_map to adjust the number of channels")
        sys.exit(1)

if not args.images_are_RGB:
    if args.use_mean_patch_embedding:
        average_conv2d_weights = torch.mean(model.patch_embed.proj.weight,1, keepdim=True)
        conv2d_weights_per_chan = average_conv2d_weights.repeat(1,len(selected_channels),1,1)
        model.patch_embed.proj.weight = nn.Parameter(conv2d_weights_per_chan)
    elif args.use_custom_embedding_map:
        embedding_seq = embedding_seq(args.custom_embedding_map, args.selected_channels)
        model = build_weight_emb(embedding_seq, model)

def prepare_img(image, args):
    if not args.images_are_RGB:
        image = utils.normalize_numpy_0_to_1(image)
        image = torch.from_numpy(image).permute(2, 0, 1)
        img = image.type(torch.FloatTensor)
        transform = transforms.Compose([])
        if args.center_crop:
            transform.transforms.append(transforms.CenterCrop(args.center_crop))
        if args.resize:
            transform.transforms.append(transforms.Resize(args.resize_length))
        if args.normalize:
            selected_channels = list(map(int, args.selected_channels))
            norm_per_channel = ast.literal_eval(args.norm_per_channel)
            mean_for_selected_channel, std_for_selected_channel = tuple([norm_per_channel[0][mean] for mean in selected_channels]), tuple([norm_per_channel[1][mean] for mean in selected_channels])
            transform.transforms.append(transforms.Normalize(mean=mean_for_selected_channel, std=std_for_selected_channel))
        img = transform(img)
    else:
        transform = transforms.Compose([
        transforms.Resize(256, interpolation=3),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ])
        img = transform(img)
    return img

def prepare_og_image(image, args, channel):
    image = utils.normalize_numpy_0_to_1(image)
    image = torch.from_numpy(image).permute(2, 0, 1)
    img = image.type(torch.FloatTensor)
    transform = transforms.Compose([])
    if args.center_crop:
        transform.transforms.append(transforms.CenterCrop(args.center_crop))
    if args.resize:
        transform.transforms.append(transforms.Resize(args.resize_length))
    if args.normalize:
        selected_channels = [channel for i in range(3)]
        norm_per_channel = ast.literal_eval(args.norm_per_channel)
        mean_for_selected_channel, std_for_selected_channel = tuple([norm_per_channel[0][mean] for mean in selected_channels]), tuple([norm_per_channel[1][mean] for mean in selected_channels])
        transform.transforms.append(transforms.Normalize(mean=mean_for_selected_channel, std=std_for_selected_channel))
    img = transform(img)
    return img

dataset_total =  datasets.ImageFolder(os.path.join(args.dataset_dir))

classes = utils.fetch_all_folder_names_of_folder_depth(dataset_total, folder_depth=args.folder_depth_for_labels)

def get_channel_name_combi(channel_combi, channel_dict):
            name_of_channel_combi = ""
            for channel_number in iter(str(channel_combi)):
                name_of_channel_combi = "_".join([name_of_channel_combi, channel_dict[int(channel_number)]])
            return name_of_channel_combi

selected_channel_str = "".join(str(x) for x in args.selected_channels)
channel_names = get_channel_name_combi(selected_channel_str, args.channel_dict)

#create directory for attention images
if args.scDINO_full_pipeline:
    checkpoint = args.pretrained_weights.split('/')[-1]
    epoch_num = re.findall(r'\d+', checkpoint)[0]
    run_directory = os.path.join(args.output_dir,args.name_of_run, f"{args.full_ViT_name}_analyses/attention_images/epoch{epoch_num}")
else:
    checkpoint = args.pretrained_weights.split('/')[-1].split('.')[0]
    run_directory = os.path.join(args.output_dir,args.name_of_run, f"attention_images/channel{channel_names}_model_{checkpoint}")
try:
    os.mkdir(run_directory)
except:
    pass
for class_name in classes:
    try:
        os.mkdir(f"{run_directory}/{class_name}")
    except:
        pass

# classes_dict = dataset_total.class_to_idx
# random.seed(args.seed)

# classes_dict = {class_name: i for i, class_name in enumerate(classes)}

print('Computing and saving attention images...')
for class_name in classes:
# for class_name, class_number in classes_dict.items():
    # class_torch = torch.tensor([class_number])
    class_indices = [i for i, x in enumerate(dataset_total.imgs) if class_name == x[0].split('/')[-args.folder_depth_for_labels-2]]
    print(f"Computing attention images for class {class_name} with {len(class_indices)} images")
    print(class_indices)
    if args.scDINO_full_pipeline: #subsample from val_indices
        validation_split = float(1-args.train_datasetsplit_fraction)
        shuffle_dataset = True
        # Creating data indices for training and validation splits:
        dataset_size = len(dataset_total)
        indices = list(range(dataset_size))
        split = int(np.floor(validation_split * dataset_size))
        if shuffle_dataset :
            np.random.seed(args.seed)
            np.random.shuffle(indices)
        train_indices, val_indices = indices[split:], indices[:split]
        val_indices_intersect = list(set(val_indices) & set(list(map(int, class_indices))))
        random_indices = random.sample(val_indices_intersect, args.num_images_per_class)
    else: #sample from whole dataset
        random_indices = random.sample(list(class_indices), args.num_images_per_class)
    for cell in random_indices:
        image_path = dataset_total.samples[cell][0]
        cell_name = image_path.split('/')[-1]
        random_cell_dir= f"{run_directory}/{class_name}/{cell_name}"
        try:
            os.mkdir(random_cell_dir)
        except:
            pass
        image_raw = imread(image_path)
        image_raw=image_raw.astype(float)
        image = image_raw[:,:,selected_channels]
        img = prepare_img(image, args)

        images_to_visualise = []
        images_to_visualise.append(img)

        if args.resize_attention_image:
            resize= torchvision.transforms.Resize((args.image_size))
            img_resized = resize(img)
            images_to_visualise.append(img_resized)
        
        #save original images
        for channel in selected_channels:
            channel_name = get_channel_name_combi(str(channel), args.channel_dict)
            multiplied_channels = [channel for i in range(3)]
            image = image_raw[:,:,multiplied_channels]
            og_img= prepare_og_image(image, args, channel)
            torchvision.utils.save_image(torchvision.utils.make_grid(og_img, normalize=True, scale_each=True), os.path.join(f"{random_cell_dir}/original_img_channel"+str(channel_name)+".png"))

        for img in images_to_visualise:
            # make the image divisible by the patch size
            w, h = img.shape[1] - img.shape[1] % args.patch_size, img.shape[2] - img.shape[2] % args.patch_size
            img = img[:, :w, :h].unsqueeze(0)

            w_featmap = img.shape[-2] // args.patch_size
            h_featmap = img.shape[-1] // args.patch_size

            try:
                attentions = model.get_last_selfattention(img.to(device))
            ## print the erorr message
            except Exception as e:
                print("Error in getting attention images for image.shape: ", img.shape)
                print(e)
                continue

            nh = attentions.shape[1] # number of head

            # we keep only the output patch attention
            attentions = attentions[0, :, 0, 1:].reshape(nh, -1)

            if args.threshold is not None:
                # we keep only a certain percentage of the mass
                val, idx = torch.sort(attentions)
                val /= torch.sum(val, dim=1, keepdim=True)
                cumval = torch.cumsum(val, dim=1)
                th_attn = cumval > (1 - args.threshold)
                idx2 = torch.argsort(idx)
                for head in range(nh):
                    th_attn[head] = th_attn[head][idx2[head]]
                th_attn = th_attn.reshape(nh, w_featmap, h_featmap).float()
                # interpolate
                th_attn = nn.functional.interpolate(th_attn.unsqueeze(0), scale_factor=args.patch_size, mode="nearest")[0].cpu().detach().numpy()

            attentions = attentions.reshape(nh, w_featmap, h_featmap)
            attentions = nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=args.patch_size, mode="nearest")[0].cpu().detach().numpy()
            
            image_size = img.shape[-1]
            try:
                final_dir = f"{random_cell_dir}/image_size_{image_size}"
                os.mkdir(final_dir)
            except:
                pass

            for j in range(nh):
                fname = os.path.join(final_dir, "attn-head" + str(j)+".png")
                plt.imsave(fname=fname, arr=attentions[j], format='png')
                

            if args.threshold is not None:
                image = skimage.io.imread(os.path.join(random_cell_dir,".png"))
        
                for j in range(nh):
                    display_instances(image, th_attn[j], fname=os.path.join(random_cell_dir, "mask_th" + str(args.threshold) + "_head" + str(j)+".png"), blur=False)
            #free gpu memory
            del img
            del attentions
    
#creating log file
with open(os.path.join(run_directory, "run_log.txt"), "w") as f:
    f.write(f"successfully computed attention visualisation with following parameters and random seed {args.seed}: \n")
    f.write("parameters: \n")
    for arg in vars(args):
        f.write(f"{arg} : {getattr(args, arg)} \n")
3.test_cases/4.DDP/pyscripts/visualise_attention.py
aws-samples/awsome-distributed-training
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

import torch

try:
    from rich.console import Console

    print = Console(force_terminal=True, force_jupyter=False).out
except ModuleNotFoundError:
    pass

print(f"{torch.cuda.is_available()=}")
print(f"{torch.backends.cuda.is_built()=}")
print(f"{torch.backends.cuda.matmul.allow_tf32=}")
print(f"{torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction=}")
print(f"{torch.backends.cuda.cufft_plan_cache=}")
print(f"{torch.backends.cuda.preferred_linalg_library(backend=None)=}")
print(f"{torch.backends.cuda.flash_sdp_enabled()=}")
print(f"{torch.backends.cuda.math_sdp_enabled()=}")

print(f"{torch.backends.cudnn.version()=}")
print(f"{torch.backends.cudnn.is_available()=}")
print(f"{torch.backends.cudnn.enabled=}")
print(f"{torch.backends.cudnn.allow_tf32=}")
print(f"{torch.backends.cudnn.deterministic=}")
print(f"{torch.backends.cudnn.benchmark=}")
print(f"{torch.backends.cudnn.benchmark_limit=}")

print(f"{torch.backends.mkl.is_available()=}")
print(f"{torch.backends.mkldnn.is_available()=}")

print(f"{torch.backends.openmp.is_available()=}")
print(f"{torch.backends.opt_einsum.is_available()=}")
print(f"{torch.backends.opt_einsum.get_opt_einsum()=}")
print(f"{torch.backends.opt_einsum.enabled=}")
print(f"{torch.backends.opt_einsum.strategy=}")

print(f"{torch.distributed.is_available()=}")
print(f"{torch.distributed.is_mpi_available()=}")
print(f"{torch.distributed.is_nccl_available()=}")

4.validation_scripts/1.pytorch-env-validation/pytorch-screen.py
aws-samples/awsome-distributed-training