25760,31136
13664,12320
37184,26656
13440,15456
37184,24416
21056,28672
24192,21728
10752,11648
32256,22176
34272,19264
29120,20832
...

A_1 = [
    'SLIDE_000',
        ...    ,
    'SLIDE_009',
]
       ...
A_5 = [
    'SLIDE_040',
        ...    ,
    'SLIDE_049',
]

B_1 = [
    'SLIDE_100',
        ...    ,
    'SLIDE_109',
]
        ...
C_5 = [
    'SLIDE_240',
        ...    ,
    'SLIDE_249',
]

# slideを訓練用とテスト(valid)用に分割
def slide_split(train, test):
    # ex) train = '123', test_or_valid = '4'

    data_map = {}
    data_map['data1'] = [A_1, B_1, C_1]
    data_map['data2'] = [A_2, B_2, C_2]
    data_map['data3'] = [A_3, B_3, C_3]
    data_map['data4'] = [A_4, B_4, C_4]
    data_map['data5'] = [A_5, B_5, C_5]

    train_list = [i for i in train]
    train_A = []
    train_B = []
    train_C = []
    for num in train_list:
        train_A = train_A + data_map[f'data{num}'][0]
        train_B = train_B + data_map[f'data{num}'][1]
        train_C = train_C + data_map[f'data{num}'][2]

    test_list = [i for i in test]
    test_A = []
    test_B = []
    test_C = []
    for num in test_list:
        test_A = test_A + data_map[f'data{num}'][0]
        test_B = test_B + data_map[f'data{num}'][1]
        test_C = test_C + data_map[f'data{num}'][2]

    return train_A, train_B, train_C, test_A, test_B, test_C

import torch
import random
import numpy as np
import os
import openslide

# map vips formats to np dtypes
format_to_dtype = {
    'uchar': np.uint8,
    'char': np.int8,
    'ushort': np.uint16,
    'short': np.int16,
    'uint': np.uint32,
    'int': np.int32,
    'float': np.float32,
    'double': np.float64,
    'complex': np.complex64,
    'dpcomplex': np.complex128,
}

DATA_PATH = f'data_directry'

class OriginalDataset(torch.utils.data.Dataset):
    def __init__(self, dataset, transform = None, bag_num=50, bag_size=100, train=True):
        self.transform = transform

        # Bagリスト作成
        # ここでは座標情報のみを扱い, 画像は読み込まない
        self.bag_list = []
        for slide_data in dataset:
            slideID = slide_data[0]
            label = slide_data[1]
            # CSVファイルからpatchの座標リスト取得
            pos = np.loadtxt(f'{DATA_PATH}/csv/{slideID}.csv', delimiter=',', dtype='int')
            if not train:
                np.random.seed(seed=int(slideID.replace('SLIDE_','')))
            np.random.shuffle(pos)
            if pos.shape[0] > bag_num*bag_size:
                pos = pos[0:(bag_num*bag_size),:]
                for i in range(bag_num):
                    patches = pos[i*bag_size:(i+1)*bag_size,:].tolist()
                    self.bag_list.append([patches, slideID, label])
            else:
                for i in range(pos.shape[0]//bag_size):
                    patches = pos[i*bag_size:(i+1)*bag_size,:].tolist()
                    self.bag_list.append([patches, slideID, label])

        random.shuffle(self.bag_list)
        self.data_num = len(self.bag_list)

    def __len__(self):
        return self.data_num

    def __getitem__(self, idx):
        pos_list = self.bag_list[idx][0]
        patch_len = len(pos_list)
        b_size = 224
        svs_list = os.listdir(f'{DATA_PATH}/svs')
        svs_fn = [s for s in svs_list if self.bag_list[idx][1] in s]
        svs = openslide.OpenSlide(f'{DATA_PATH}/svs/{svs_fn[0]}')
        bag = torch.empty(patch_len, 3, 224, 224, dtype=torch.float)

        # 画像読み込み
        i = 0
        for pos in pos_list:
            if self.transform:
                img = svs.read_region((pos[0],pos[1]),0,(b_size,b_size)).convert('RGB')
                img = self.transform(img)
                bag[i] = img
            i += 1

        label = self.bag_list[idx][2]
        label = torch.LongTensor([label])

        # バッグ, ラベルを返す
        return bag, label

from torch.nn.parallel import DistributedDataParallel as DDP

model = model.to(rank)
process_group = torch.distributed.new_group([i for i in range(world_size)])
# modelのBatchNormをSyncBatchNormに変更してくれる
model = nn.SyncBatchNorm.convert_sync_batchnorm(model, process_group)
# modelをmulti GPU対応させる
ddp_model = DDP(model, device_ids=[rank])

# Train bag作成(epochごとにbag再構築)
original_dataset = OriginalDataset(
    transform=transform,
    dataset=train_dataset,
    train=True
)

# Datasetをmulti GPU対応させる
deta_sampler = torch.utils.data.distributed.DistributedSampler(original_dataset, rank=rank)

# batch_sizeで設定した個数だけbagを各GPUに分配
train_loader = torch.utils.data.DataLoader(
    original_dataset,
    batch_size=1,
    shuffle=False,
    pin_memory=False,
    num_workers=2,
    sampler=deta_sampler
)

    # model読み込み
    from MIL import feature_extractor, class_predictor, MIL

    # 特徴抽出器の作成 (今回はResNet50を使用)
    encoder = models.resnet50(pretrained=True)
    encoder.fc = nn.Identity()
    for p in encoder.parameters():
        p.required_grad = True

    # MILモデルの構築
    feature_ex = feature_extractor(encoder)
    class_pred = class_predictor()
    model = MIL(feature_ex, class_pred)
    model = model.to(rank)
    process_group = torch.distributed.new_group([i for i in range(world_size)])
    # modelのBatchNormをSyncBatchNormに変更してくれる
    model = nn.SyncBatchNorm.convert_sync_batchnorm(model, process_group)
    # modelをmulti GPU対応させる
    ddp_model = DDP(model, device_ids=[rank])

    # クロスエントロピー損失関数使用
    loss_fn = nn.CrossEntropyLoss()
    # SGDmomentum法使用
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001)

    # 前処理
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((224, 224)),
        torchvision.transforms.ToTensor()
    ])

    # Train
    for epoch in range(EPOCHS):
        # Train bag作成(epochごとにbag再構築)
        original_dataset = OriginalDataset(
            transform=transform,
            dataset=train_dataset,
            train=True
        )

        # Datasetをmulti GPU対応させる
        deta_sampler = torch.utils.data.distributed.DistributedSampler(original_dataset, rank=rank)

        # batch_sizeで設定した個数だけbagを各GPUに分配
        train_loader = torch.utils.data.DataLoader(
            original_dataset,
            batch_size=1,
            shuffle=False,
            pin_memory=False,
            num_workers=2,
            sampler=deta_sampler
        )

        # 学習
        class_loss, acc = train(ddp_model, rank, loss_fn, optimizer, train_loader)

def train(model, rank, loss_fn, optimizer, train_loader):
    model.train()
    train_class_loss = 0.0
    correct_num = 0

    for (input_tensor, class_label) in train_loader:
        # Bagとbatchのギャップを吸収して学習
        for bag_num in range(input_tensor.shape[0]):
            bag = input_tensor[bag_num].to(rank, non_blocking=True)
            bag_class_label = class_label[bag_num].to(rank, non_blocking=True)
            class_prob, class_hat = model(bag)
            class_loss = loss_fn(class_prob, bag_class_label)
            train_class_loss += class_loss.item()
            optimizer.zero_grad()
            class_loss.backward()
            optimizer.step() 
            correct_num += eval_ans(class_hat, bag_class_label)

    train_class_loss = train_class_loss / len(train_loader)
    train_acc = correct_num / len(train_loader)

    return train_class_loss, train_acc