code: import logging
  
  import os
  import numpy as np
  import torch
  from torch.utils.data import Dataset, DataLoader
  import torchvision.transforms as transforms
  from torch.utils.data.distributed import DistributedSampler
  
  from .dataset import CheXpert
  
  
  def _get_mean_and_std(dataset: Dataset):
      """Compute the mean and std of dataset."""
      data_loader = DataLoader(dataset, batch_size=1, shuffle=False)
      mean = torch.zeros(3)
      std = torch.zeros(3)
      for i, (img, _) in enumerate(data_loader):
          if i % 1000 == 0:
              print(i)
          mean += img.mean(dim=(0, 2, 3))
          std += img.std(dim=(0, 2, 3))
      mean /= len(data_loader)
      std /= len(data_loader)
      return mean, std
  
  
  class Cutout(object):
      def __init__(self, length):
          self.length = length
  
      def __call__(self, img):
          h, w = img.size(1), img.size(2)
          mask = np.ones((h, w), np.float32)
          y = np.random.randint(h)
          x = np.random.randint(w)
  
          y1 = np.clip(y - self.length // 2, 0, h)
          y2 = np.clip(y + self.length // 2, 0, h)
          x1 = np.clip(x - self.length // 2, 0, w)
          x2 = np.clip(x + self.length // 2, 0, w)
  
          mask[y1:y2, x1:x2] = 0.0
          mask = torch.from_numpy(mask)
          mask = mask.expand_as(img)
          img *= mask
          return img
  
  
  def _data_transforms_chexpert():
  
      CHEXPERT_MEAN = [0.503, 0.503, 0.503]
      CHEXPERT_STD = [0.291, 0.291, 0.291]
  
      image_size = 256
      train_transform = transforms.Compose(
          [
              # transforms.ToPILImage(),
              transforms.RandomResizedCrop(image_size),
              # transforms.RandomHorizontalFlip(),
              transforms.ToTensor(),
              transforms.Normalize(CHEXPERT_MEAN, CHEXPERT_STD),
          ]
      )
  
      # train_transform.transforms.append(Cutout(16))
  
      test_transform = transforms.Compose(
          [
              transforms.CenterCrop(image_size),
              transforms.ToTensor(),
              transforms.Normalize(CHEXPERT_MEAN, CHEXPERT_STD),
          ]
      )
  
      return train_transform, test_transform
  
  
  # for centralized training
  def get_dataloader(dataset, datadir, train_bs, test_bs, dataidxs=None, policy="zeros"):
      return get_dataloader_chexpert(datadir, train_bs, test_bs, dataidxs, policy=policy)
  
  
  # for local devices
  def get_dataloader_test(dataset, datadir, train_bs, test_bs, dataidxs_train, dataidxs_test, policy="zeros"):
      return get_dataloader_test_chexpert(datadir, train_bs, test_bs, dataidxs_train, dataidxs_test, policy=policy)
  
  
  def get_dataloader_chexpert(datadir, train_bs, test_bs, dataidxs=None, policy="zeros"):
      dl_obj = CheXpert
  
      transform_train, transform_test = _data_transforms_chexpert()
  
      train_ds = dl_obj(
          datadir,
          dataidxs=dataidxs,
          train=True,
          transform=transform_train,
          download=False,
          policy=policy,
      )
      test_ds = dl_obj(
          datadir,
          dataidxs=None,
          train=False,
          transform=transform_test,
          download=False,
          policy=policy,
      )
  
      train_dl = DataLoader(
          dataset=train_ds,
          batch_size=train_bs,
          shuffle=True,
          drop_last=False,
          pin_memory=True,
          num_workers=4,
      )
      test_dl = DataLoader(
          dataset=test_ds,
          batch_size=test_bs,
          shuffle=False,
          drop_last=False,
          pin_memory=True,
          num_workers=4,
      )
  
      return train_dl, test_dl
  
  
  def get_dataloader_test_chexpert(datadir, train_bs, test_bs, dataidxs_train=None, dataidxs_test=None, policy="zeros"):
      dl_obj = CheXpert
  
      transform_train, transform_test = _data_transforms_chexpert()
  
      train_ds = dl_obj(
          datadir,
          dataidxs=dataidxs_train,
          train=True,
          transform=transform_train,
          download=True,
          policy=policy,
      )
      test_ds = dl_obj(
          datadir,
          dataidxs=dataidxs_test,
          train=False,
          transform=transform_test,
          download=True,
          policy=policy,
      )
  
      train_dl = DataLoader(
          dataset=train_ds,
          batch_size=train_bs,
          shuffle=True,
          drop_last=False,
          pin_memory=True,
          num_workers=4,
      )
      test_dl = DataLoader(
          dataset=test_ds,
          batch_size=test_bs,
          shuffle=False,
          drop_last=False,
          pin_memory=True,
          num_workers=4,
      )
  
      return train_dl, test_dl
  
  
  def distributed_centralized_chexpert_loader(dataset, data_dir, world_size, rank, batch_size):
      """
      Used for generating distributed dataloader for
      accelerating centralized training
      """
  
      train_bs = batch_size
      test_bs = batch_size
  
      transform_train, transform_test = _data_transforms_chexpert()
      train_dataset = CheXpert(data_dir=data_dir, dataidxs=None, train=True, transform=transform_train)
      test_dataset = CheXpert(data_dir=data_dir, dataidxs=None, train=False, transform=transform_test)
  
      train_sam = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
      test_sam = DistributedSampler(test_dataset, num_replicas=world_size, rank=rank)
  
      train_dl = data.DataLoader(
          train_dataset,
          batch_size=train_bs,
          sampler=train_sam,
          pin_memory=True,
          num_workers=4,
      )
      test_dl = data.DataLoader(
          test_dataset,
          batch_size=test_bs,
          sampler=test_sam,
          pin_memory=True,
          num_workers=4,
      )
  
      class_num = 1000
  
      train_data_num = len(train_dataset)
      test_data_num = len(test_dataset)
  
      return train_data_num, test_data_num, train_dl, test_dl, None, None, None, class_num
  
  
  def load_partition_data_chexpert(
      data_dir,
      partition_method="random",
      partition_alpha=None,
      client_number=100,
      batch_size=10,
      policy="zeros",
  ):
      transform_train, transform_test = _data_transforms_chexpert()
  
      train_dataset = CheXpert(
          data_dir=data_dir,
          dataidxs=None,
          train=True,
          transform=transform_train,
          policy=policy,
      )
      test_dataset = CheXpert(data_dir=data_dir, dataidxs=None, train=False, transform=transform_test, policy=policy)
  
      # get local dataset
      if partition_method == "random":
          num_train_items = int(len(train_dataset) / client_number)
          num_test_items = int(len(test_dataset) / client_number)
          dict_client = {}
          all_train_idxs = list(range(len(train_dataset)))
          all_test_idxs = list(range(len(test_dataset)))
          for client_idx in range(client_number):
              dict_client[client_idx] = {}
              dict_client[client_idx]["train"] = set(np.random.choice(all_train_idxs, num_train_items, replace=False))
              dict_client[client_idx]["test"] = set(np.random.choice(all_test_idxs, num_test_items, replace=False))
              all_train_idxs = list(set(all_train_idxs) - dict_client[client_idx]["train"])
              all_test_idxs = list(set(all_test_idxs) - dict_client[client_idx]["test"])
          if len(all_train_idxs) > 0:
              all_client_idxs = list(range(client_number))
              np.random.shuffle(all_client_idxs)
              choiced_client_idxs = all_client_idxs[: len(all_train_idxs)]
              for idx, client_idx in enumerate(choiced_client_idxs):
                  dict_client[client_idx]["train"].add(all_train_idxs[idx])
          if len(all_test_idxs) > 0:
              all_client_idxs = list(range(client_number))
              np.random.shuffle(all_client_idxs)
              choiced_client_idxs = all_client_idxs[: len(all_test_idxs)]
              for idx, client_idx in enumerate(choiced_client_idxs):
                  dict_client[client_idx]["test"].add(all_test_idxs[idx])
      else:
          raise NotImplementedError
  
      # build dataloader
      train_dl = []
      test_dl = []
      for client_idx in range(client_number):
          train_data_idxs = list(dict_client[client_idx]["train"])
          test_data_idxs = list(dict_client[client_idx]["test"])
          train_dl_, test_dl_ = get_dataloader_test_chexpert(
              datadir=data_dir,
              dataidxs_train=train_data_idxs,
              dataidxs_test=test_data_idxs,
              train_bs=batch_size,
              test_bs=batch_size,
              policy=policy,
          )
          train_dl.append(train_dl_)
          test_dl.append(test_dl_)
  
          logging.info(f"Client {client_idx} train data num: {len(train_dl_)} test data num: {len(test_dl_)}")
  
      logging.info("Partition data done")
      # logging.info("Partition data for each client: {}".format(dict_client))
  
      train_data_num = len(train_dataset)
      test_data_num = len(test_dataset)
      train_data_global = train_dataset
      test_data_global = test_dataset
      data_local_num_dict = {
          client_idx: len(dict_client[client_idx]["train"]) + len(dict_client[client_idx]["test"])
          for client_idx in range(client_number)
      }
      train_data_local_dict = {client_idx: train_dl_ for client_idx, train_dl_ in enumerate(train_dl)}
      test_data_local_dict = {client_idx: test_dl_ for client_idx, test_dl_ in enumerate(test_dl)}
      class_num = train_dataset.num_classes
  
      return (
          train_data_num,
          test_data_num,
          train_data_global,
          test_data_global,
          data_local_num_dict,
          train_data_local_dict,
          test_data_local_dict,
          class_num,
      )
  
  
  if __name__ == "__main__":
      data_path = os.path.join("D:\\", "dataset", "CheXpert", "CheXpert-v1.0-small")
      data = CheXpert(data_dir=data_path, transform=transforms.ToTensor())
      print(len(data))
      print(data[0][0])
      print(data[0][1])
  
      # mean, std = _get_mean_and_std(data)
      # print(mean, std)
  
      # train_transform, valid_transform = _data_transforms_chexpert()
      # print(train_transform)
      # print(valid_transform)
  
      (
          train_data_num,
          test_data_num,
          train_data_global,
          test_data_global,
          data_local_num_dict,
          train_data_local_dict,
          test_data_local_dict,
          class_num,
      ) = load_partition_data_chexpert(data_dir=data_path, client_number=10, batch_size=10, policy="zeros")
  
      print(train_data_num, test_data_num, class_num)
   import numpy as np
  import scipy.sparse as sp
  import torch
  from sklearn.preprocessing import StandardScaler
  from sklearn.model_selection import train_test_split
  
  from torch_geometric.utils import to_networkx, degree
  import torch.nn.functional as F
  
  
  def convert_to_nodeDegreeFeatures(graphs):
      # print(graph.x)
      graph_infos = []
      maxdegree = 0
      for i, graph in enumerate(graphs):
          g = to_networkx(graph, to_undirected=True)
          gdegree = max(dict(g.degree).values())
          if gdegree > maxdegree:
              maxdegree = gdegree
          graph_infos.append(
              (graph, g.degree, graph.num_nodes)
          )  # (graph, node_degrees, num_nodes)
  
      new_graphs = []
      for i, tuple in enumerate(graph_infos):
          idx, x = tuple[0].edge_index[0], tuple[0].x
          deg = degree(idx, tuple[2], dtype=torch.long)
          deg = F.one_hot(deg, num_classes=maxdegree + 1).to(torch.float)
  
          new_graph = tuple[0].clone()
          new_graph.__setitem__("x", deg)
          new_graphs.append(new_graph)
  
      return new_graphs
  
  
  def split_data(graphs, train=None, test=None, shuffle=True, seed=None):
      y = torch.cat([graph.y for graph in graphs])
      graphs_tv, graphs_test = train_test_split(
          graphs,
          train_size=train,
          test_size=test,
          stratify=y,
          shuffle=shuffle,
          random_state=seed,
      )
      return graphs_tv, graphs_test
  
  
  def np_uniform_sample_next(compact_adj, tree, fanout):
      last_level = tree[-1]  # [batch, f^depth]
      batch_lengths = compact_adj.degrees[last_level]
      nodes = np.repeat(last_level, fanout, axis=1)
      batch_lengths = np.repeat(batch_lengths, fanout, axis=1)
      batch_next_neighbor_ids = np.random.uniform(
          size=batch_lengths.shape, low=0, high=1 - 1e-9
      )
      # Shape = (len(nodes), neighbors_per_node)
      batch_next_neighbor_ids = np.array(
          batch_next_neighbor_ids * batch_lengths, dtype=last_level.dtype
      )
      shape = batch_next_neighbor_ids.shape
      batch_next_neighbor_ids = np.array(
          compact_adj.compact_adj[nodes.reshape(-1), batch_next_neighbor_ids.reshape(-1)]
      ).reshape(shape)
  
      return batch_next_neighbor_ids
  
  
  def np_traverse(
      compact_adj, seed_nodes, fanouts=(1,), sample_fn=np_uniform_sample_next
  ):
      if not isinstance(seed_nodes, np.ndarray):
          raise ValueError("Seed must a numpy array")
  
      if (
          len(seed_nodes.shape) > 2
          or len(seed_nodes.shape) < 1
          or not str(seed_nodes.dtype).startswith("int")
      ):
          raise ValueError("seed_nodes must be 1D or 2D int array")
  
      if len(seed_nodes.shape) == 1:
          seed_nodes = np.expand_dims(seed_nodes, 1)
  
      # Make walk-tree
      forest_array = [seed_nodes]
      for f in fanouts:
          next_level = sample_fn(compact_adj, forest_array, f)
          assert next_level.shape[1] == forest_array[-1].shape[1] * f
  
          forest_array.append(next_level)
  
      return forest_array
  
  
  class WalkForestCollator(object):
      def __init__(self, normalize_features=False):
          self.normalize_features = normalize_features
  
      def __call__(self, molecule):
          comp_adj, feature_matrix, label, fanouts = molecule[0]
          node_ids = np.array(list(range(feature_matrix.shape[0])), dtype=np.int32)
          forest = np_traverse(comp_adj, node_ids, fanouts)
          torch_forest = [torch.from_numpy(forest[0]).flatten()]
          label = np.where(np.isnan(label), 0.0, label)
  
          for i in range(len(forest) - 1):
              torch_forest.append(torch.from_numpy(forest[i + 1]).reshape(-1, fanouts[i]))
  
          if self.normalize_features:
              mx = sp.csr_matrix(feature_matrix)
              rowsum = np.array(mx.sum(1))
              r_inv = np.power(rowsum, -1).flatten()
              r_inv[np.isinf(r_inv)] = 0.0
              r_mat_inv = sp.diags(r_inv)
              normalized_feature_matrix = r_mat_inv.dot(mx)
              normalized_feature_matrix = np.array(normalized_feature_matrix.todense())
          else:
              scaler = StandardScaler()
              scaler.fit(feature_matrix)
              normalized_feature_matrix = scaler.transform(feature_matrix)
  
          return (
              torch_forest,
              torch.as_tensor(normalized_feature_matrix, dtype=torch.float32),
              torch.as_tensor(label, dtype=torch.float32),
          )
  
  
  class DefaultCollator(object):
      def __init__(self, normalize_features=True, normalize_adj=True):
          self.normalize_features = normalize_features
          self.normalize_adj = normalize_adj
  
      def __call__(self, molecule):
          adj_matrix, feature_matrix, label, _ = molecule[0]
          label = np.where(np.isnan(label), 0.0, label)
  
          if self.normalize_features:
              mx = sp.csr_matrix(feature_matrix)
              rowsum = np.array(mx.sum(1))
              r_inv = np.power(rowsum, -1).flatten()
              r_inv[np.isinf(r_inv)] = 0.0
              r_mat_inv = sp.diags(r_inv)
              normalized_feature_matrix = r_mat_inv.dot(mx)
              normalized_feature_matrix = np.array(normalized_feature_matrix.todense())
          else:
              scaler = StandardScaler()
              scaler.fit(feature_matrix)
              normalized_feature_matrix = scaler.transform(feature_matrix)
  
          if self.normalize_adj:
              rowsum = np.array(adj_matrix.sum(1))
              r_inv_sqrt = np.power(rowsum, -0.5).flatten()
              r_inv_sqrt[np.isinf(r_inv_sqrt)] = 0.0
              r_mat_inv_sqrt = sp.diags(r_inv_sqrt)
              normalized_adj_matrix = (
                  adj_matrix.dot(r_mat_inv_sqrt).transpose().dot(r_mat_inv_sqrt)
              )
          else:
              normalized_adj_matrix = adj_matrix
  
          return (
              torch.as_tensor(
                  np.array(normalized_adj_matrix.todense()), dtype=torch.float32
              ),
              torch.as_tensor(normalized_feature_matrix, dtype=torch.float32),
              torch.as_tensor(label, dtype=torch.float32),
          )
apis: torch.utils.data.DataLoader numpy.ones numpy.random.shuffle torch.utils.data.distributed.DistributedSampler numpy.random.choice numpy.clip torch.from_numpy torch.zeros numpy.random.randint
apis: numpy.random.uniform torch.as_tensor numpy.isinf scipy.sparse.csr_matrix numpy.repeat scipy.sparse.diags torch.nn.functional.one_hot numpy.expand_dims torch.from_numpy numpy.power numpy.isnan numpy.array torch.cat sklearn.model_selection.train_test_split sklearn.preprocessing.StandardScaler