json文件：
{
"pathmnist": {
"description": "PathMNIST: A dataset based on a prior study for predicting survival from colorectal cancer histology slides, which provides a dataset NCT-CRC-HE-100K of 100,000 non-overlapping image patches from hematoxylin & eosin stained histological images, and a test dataset CRC-VAL-HE-7K of 7,180 image patches from a different clinical center. 9 types of tissues are involved, resulting a multi-class classification task. We resize the source images of 3 x 224 x 224 into 3 x 28 x 28, and split NCT-CRC-HE-100K into training and valiation set with a ratio of 9:1.",
"url": "https://zenodo.org/record/4269852/files/pathmnist.npz?download=1",
"MD5": "a8b06965200029087d5bd730944a56c1",
"task": "multi-class",
"label": {
"0": "adipose",
"1": "background",
"2": "debris",
"3": "lymphocytes",
"4": "mucus",
"5": "smooth muscle",
"6": "normal colon mucosa",
"7": "cancer-associated stroma",
"8": "colorectal adenocarcinoma epithelium"
},
"n_channels": 3,
"n_samples": {
"train": 89996,
"val": 10004,
"test": 7180
},
"license": "CC BY 4.0"
},
"chestmnist": {
"description": "ChestMNIST: A dataset based on NIH-ChestXray14 dataset, comprising 112,120 frontal-view X-ray images of 30,805 unique patients with the text-mined 14 disease image labels, which could be formulized as multi-label binary classification task. We use the official data split, and resize the source images of 1 x 1024 x 1024 into 1 x 28 x 28.",
"url": "https://zenodo.org/record/4269852/files/chestmnist.npz?download=1",
"MD5": "02c8a6516a18b556561a56cbdd36c4a8",
"task": "multi-label, binary-class",
"label": {
"0": "atelectasis",
"1": "cardiomegaly",
"2": "effusion",
"3": "infiltration",
"4": "mass",
"5": "nodule",
"6": "pneumonia",
"7": "pneumothorax",
"8": "consolidation",
"9": "edema",
"10": "emphysema",
"11": "fibrosis",
"12": "pleural",
"13": "hernia"
},
"n_channels": 1,
"n_samples": {
"train": 78468,
"val": 11219,
"test": 22433
},
"license": "CC0 1.0"
},
"dermamnist": {
"description": "DermaMNIST: A dataset based on HAM10000, a large collection of multi-source dermatoscopic images of common pigmented skin lesions. The dataset consists of 10,015 dermatoscopic images labeled as 7 different categories, as a multi-class classification task. We split the images into training, validation and test set with a ratio of 7:1:2. The source images of 3 x 600 x 450 are resized into 3 x 28 x 28.",
"url": "https://zenodo.org/record/4269852/files/dermamnist.npz?download=1",
"MD5": "0744692d530f8e62ec473284d019b0c7",
"task": "multi-class",
"label": {
"0": "actinic keratoses and intraepithelial carcinoma",
"1": "basal cell carcinoma",
"2": "benign keratosis-like lesions",
"3": "dermatofibroma",
"4": "melanoma",
"5": "melanocytic nevi",
"6": "vascular lesions"
},
"n_channels": 3,
"n_samples": {
"train": 7007,
"val": 1003,
"test": 2005
},
"license": "CC BY-NC 4.0"
},
"octmnist": {
"description": "OCTMNIST: A dataset based on a prior dataset of 109,309 valid optical coherence tomography (OCT) images for retinal diseases. 4 types are involved, leading to a multi-class classification task. We split the source training set with a ratio of 9:1 into training and validation set, and use its source validation set as the test set. The source images are single-channel, and their sizes range from (384-1,536) x (277-512). We center-crop the images and resize them into 1 x 28 x 28.",
"url": "https://zenodo.org/record/4269852/files/octmnist.npz?download=1",
"MD5": "c68d92d5b585d8d81f7112f81e2d0842",
"task": "multi-class",
"label": {
"0": "choroidal neovascularization",
"1": "diabetic macular edema",
"2": "drusen",
"3": "normal"
},
"n_channels": 1,
"n_samples": {
"train": 97477,
"val": 10832,
"test": 1000
},
"license": "CC BY 4.0"
},
"pneumoniamnist": {
"description": "PneumoniaMNIST: A dataset based on a prior dataset of 5,856 pediatric chest X-ray images. The task is binary-class classification of pneumonia and normal. We split the source training set with a ratio of 9:1 into training and validation set, and use its source validation set as the test set. The source images are single-channel, and their sizes range from (384-2,916) x (127-2,713). We center-crop the images and resize them into 1 x 28 x 28.",
"url": "https://zenodo.org/record/4269852/files/pneumoniamnist.npz?download=1",
"MD5": "28209eda62fecd6e6a2d98b1501bb15f",
"task": "binary-class",
"label": {
"0": "normal",
"1": "pneumonia"
},
"n_channels": 1,
"n_samples": {
"train": 4708,
"val": 524,
"test": 624
},
"license": "CC BY 4.0"
},
"retinamnist": {
"description": "RetinaMNIST: A dataset based on DeepDRiD, a dataset of 1,600 retina fundus images. The task is ordinal regression for 5-level grading of diabetic retinopathy severity. We split the source training set with a ratio of 9:1 into training and validation set, and use the source validation set as test set. The source images of 3 x 1,736 x 1,824 are center-cropped and resized into 3 x 28 x 28",
"url": "https://zenodo.org/record/4269852/files/retinamnist.npz?download=1",
"MD5": "bd4c0672f1bba3e3a89f0e4e876791e4",
"task": "ordinal regression",
"label": {
"0": "0",
"1": "1",
"2": "2",
"3": "3",
"4": "4"
},
"n_channels": 3,
"n_samples": {
"train": 1080,
"val": 120,
"test": 400
},
"license": "CC BY 4.0"
},
"breastmnist": {
"description": "BreastMNIST: A dataset based on a dataset of 780 breast ultrasound images. It is categorized into 3 classes: normal, benign and malignant. As we use low-resolution images, we simplify the task into binary classification by combing normal and benign as negative, and classify them against malignant as positive. We split the source dataset with a ratio of 7:1:2 into training, validation and test set. The source images of 1 x 500 x 500 are resized into 1 x 28 x 28.",
"url": "https://zenodo.org/record/4269852/files/breastmnist.npz?download=1",
"MD5": "750601b1f35ba3300ea97c75c52ff8f6",
"task": "binary-class",
"label": {
"0": "malignant",
"1": "normal, benign"
},
"n_channels": 1,
"n_samples": {
"train": 546,
"val": 78,
"test": 156
},
"license": "CC BY 4.0"
},
"organmnist_axial": {
"description": "OrganMNIST_Axial: A dataset based on 3D computed tomography (CT) images from Liver Tumor Segmentation Benchmark (LiTS). We use bounding-box annotations of 11 body organs from another study to obtain the organ labels. Hounsfield-Unit (HU) of the 3D images are transformed into grey scale with a abdominal window; we then crop 2D images from the center slices of the 3D bounding boxes in axial views (planes). The images are resized into 1 x 28 x 28 to perform multi-class classification of 11 body organs. 115 and 16 CT scans from the source training set are used as training and validation set, respectively. The 70 CT scans from the source test set are treated as the test set.",
"url": "https://zenodo.org/record/4269852/files/organmnist_axial.npz?download=1",
"MD5": "866b832ed4eeba67bfb9edee1d5544e6",
"task": "multi-class",
"label": {
"0": "bladder",
"1": "femur-left",
"2": "femur-right",
"3": "heart",
"4": "kidney-left",
"5": "kidney-right",
"6": "liver",
"7": "lung-left",
"8": "lung-right",
"9": "pancreas",
"10": "spleen"
},
"n_channels": 1,
"n_samples": {
"train": 34581,
"val": 6491,
"test": 17778
},
"license": "CC BY 4.0"
},
"organmnist_coronal": {
"description": "OrganMNIST_Coronal: A dataset based on 3D computed tomography (CT) images from Liver Tumor Segmentation Benchmark (LiTS). We use bounding-box annotations of 11 body organs from another study to obtain the organ labels. Hounsfield-Unit (HU) of the 3D images are transformed into grey scale with a abdominal window; we then crop 2D images from the center slices of the 3D bounding boxes in coronal views (planes). The images are resized into 1 x 28 x 28 to perform multi-class classification of 11 body organs. 115 and 16 CT scans from the source training set are used as training and validation set, respectively. The 70 CT scans from the source test set are treated as the test set.",
"url": "https://zenodo.org/record/4269852/files/organmnist_coronal.npz?download=1",
"MD5": "0afa5834fb105f7705a7d93372119a21",
"task": "multi-class",
"label": {
"0": "bladder",
"1": "femur-left",
"2": "femur-right",
"3": "heart",
"4": "kidney-left",
"5": "kidney-right",
"6": "liver",
"7": "lung-left",
"8": "lung-right",
"9": "pancreas",
"10": "spleen"
},
"n_channels": 1,
"n_samples": {
"train": 13000,
"val": 2392,
"test": 8268
},
"license": "CC BY 4.0"
},
"organmnist_sagittal": {
"description": "OrganMNIST_Sagittal: A dataset based on 3D computed tomography (CT) images from Liver Tumor Segmentation Benchmark (LiTS). We use bounding-box annotations of 11 body organs from another study to obtain the organ labels. Hounsfield-Unit (HU) of the 3D images are transformed into grey scale with a abdominal window; we then crop 2D images from the center slices of the 3D bounding boxes in sagittal views (planes). The images are resized into 1 x 28 x 28 to perform multi-class classification of 11 body organs. 115 and 16 CT scans from the source training set are used as training and validation set, respectively. The 70 CT scans from the source test set are treated as the test set.",
"url": "https://zenodo.org/record/4269852/files/organmnist_sagittal.npz?download=1",
"MD5": "e5c39f1af030238290b9557d9503af9d",
"task": "multi-class",
"label": {
"0": "bladder",
"1": "femur-left",
"2": "femur-right",
"3": "heart",
"4": "kidney-left",
"5": "kidney-right",
"6": "liver",
"7": "lung-left",
"8": "lung-right",
"9": "pancreas",
"10": "spleen"
},
"n_channels": 1,
"n_samples": {
"train": 13940,
"val": 2452,
"test": 8829
},
"license": "CC BY 4.0"
}
}
dataset.py:
import os
import json
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
INFO = "medmnist/medmnist.json"
class MedMNIST(Dataset):
flag = ...
def __init__(self, root, split='train', transform=None, target_transform=None, download=False):
''' dataset
:param split: 'train', 'val' or 'test', select subset
:param transform: data transformation
:param target_transform: target transformation
'''
with open(INFO, 'r') as f:
self.info = json.load(f)[self.flag]
self.root = root
if download:
self.download()
if not os.path.exists(os.path.join(self.root, "{}.npz".format(self.flag))):
raise RuntimeError('Dataset not found.' +
' You can use download=True to download it')
npz_file = np.load(os.path.join(self.root, "{}.npz".format(self.flag)))
self.split = split
self.transform = transform
self.target_transform = target_transform
if self.split == 'train':
self.img = npz_file['train_images']
self.label = npz_file['train_labels']
elif self.split == 'val':
self.img = npz_file['val_images']
self.label = npz_file['val_labels']
elif self.split == 'test':
self.img = npz_file['test_images']
self.label = npz_file['test_labels']
def __getitem__(self, index):
img, target = self.img[index], self.label[index].astype(int)
img = Image.fromarray(np.uint8(img))
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target
def __len__(self):
return self.img.shape[0]
def __repr__(self):
'''Adapted from torchvision.
'''
_repr_indent = 4
head = "Dataset " + self.__class__.__name__
body = ["Number of datapoints: {}".format(self.__len__())]
body.append("Root location: {}".format(self.root))
body.append("Split: {}".format(self.split))
body.append("Task: {}".format(self.info["task"]))
body.append("Number of channels: {}".format(self.info["n_channels"]))
body.append("Meaning of labels: {}".format(self.info["label"]))
body.append("Number of samples: {}".format(self.info["n_samples"]))
body.append("Description: {}".format(self.info["description"]))
body.append("License: {}".format(self.info["license"]))
if hasattr(self, "transforms") and self.transforms is not None:
body += [repr(self.transforms)]
lines = [head] + [" " * _repr_indent + line for line in body]
return '\n'.join(lines)
def download(self):
try:
from torchvision.datasets.utils import download_url
download_url(url=self.info["url"], root=self.root,
filename="{}.npz".format(self.flag), md5=self.info["MD5"])
except:
raise RuntimeError('Something went wrong when downloading! ' +
'Go to the homepage to download manually. ' +
'https://github.com/MedMNIST/MedMNIST')
class PathMNIST(MedMNIST):
flag = "pathmnist"
class OCTMNIST(MedMNIST):
flag = "octmnist"
class PneumoniaMNIST(MedMNIST):
flag = "pneumoniamnist"
class ChestMNIST(MedMNIST):
flag = "chestmnist"
class DermaMNIST(MedMNIST):
flag = "dermamnist"
class RetinaMNIST(MedMNIST):
flag = "retinamnist"
class BreastMNIST(MedMNIST):
flag = "breastmnist"
class OrganMNISTAxial(MedMNIST):
flag = "organmnist_axial"
class OrganMNISTCoronal(MedMNIST):
flag = "organmnist_coronal"
class OrganMNISTSagittal(MedMNIST):
flag = "organmnist_sagittal"
models.py:
import torch.nn as nn
import torch.nn.functional as F
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, in_planes, planes, stride=1):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(
in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.shortcut = nn.Sequential()
if stride != 1 or in_planes != self.expansion*planes:
self.shortcut = nn.Sequential(
nn.Conv2d(in_planes, self.expansion*planes,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(self.expansion*planes)
)
def forward(self, x):
#print(x.shape)
out = F.relu(self.bn1(self.conv1(x)))
#print(out.shape)
out = self.bn2(self.conv2(out))
#print(out.shape)
out += self.shortcut(x)
#print(out.shape)
out = F.relu(out)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, in_planes, planes, stride=1):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
stride=stride, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes, self.expansion *
planes, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(self.expansion*planes)
self.shortcut = nn.Sequential()
if stride != 1 or in_planes != self.expansion*planes:
self.shortcut = nn.Sequential(
nn.Conv2d(in_planes, self.expansion*planes,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(self.expansion*planes)
)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = F.relu(self.bn2(self.conv2(out)))
out = self.bn3(self.conv3(out))
out += self.shortcut(x)
out = F.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block, num_blocks, in_channels=1, num_classes=2):
super(ResNet, self).__init__()
self.in_planes = 64
self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
self.linear = nn.Linear(512 * block.expansion, num_classes)
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1]*(num_blocks-1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return nn.Sequential(*layers)
def forward(self, x):
#print(x.shape)
out = F.relu(self.bn1(self.conv1(x)))
#print(out.shape)
out = self.layer1(out) #特征图个数没变，输入是64 输出也是64 在shortcut中不需要调整
#print(out.shape)
out = self.layer2(out)
#print(out.shape)
out = self.layer3(out)
#print(out.shape)
out = self.layer4(out)
#print(out.shape)
out = F.avg_pool2d(out, 4)
#print(out.shape)
out = out.view(out.size(0), -1)
#print(out.shape)
out = self.linear(out)
#print(out.shape)
return out
def ResNet18(in_channels, num_classes):
return ResNet(BasicBlock, [2, 2, 2, 2], in_channels=in_channels, num_classes=num_classes)
def ResNet50(in_channels, num_classes):
return ResNet(Bottleneck, [3, 4, 6, 3], in_channels=in_channels, num_classes=num_classes)
train.py:
import os
import argparse
import json
from tqdm import trange
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision.transforms as transforms
from medmnist.models import ResNet18
from medmnist.dataset import INFO, PathMNIST, ChestMNIST, DermaMNIST, OCTMNIST, PneumoniaMNIST, RetinaMNIST, \
BreastMNIST, OrganMNISTAxial, OrganMNISTCoronal, OrganMNISTSagittal
from medmnist.evaluator import getAUC, getACC, save_results
def main(flag, input_root, output_root, end_epoch, download):
''' main function
:param flag: name of subset
'''
dataclass = {
"pathmnist": PathMNIST,
"chestmnist": ChestMNIST,
"dermamnist": DermaMNIST,
"octmnist": OCTMNIST,
"pneumoniamnist": PneumoniaMNIST,
"retinamnist": RetinaMNIST,
"breastmnist": BreastMNIST,
"organmnist_axial": OrganMNISTAxial,
"organmnist_coronal": OrganMNISTCoronal,
"organmnist_sagittal": OrganMNISTSagittal,
}
with open(INFO, 'r') as f:
info = json.load(f)
task = info[flag]['task']
n_channels = info[flag]['n_channels']
n_classes = len(info[flag]['label'])
start_epoch = 0
lr = 0.001
batch_size = 128
val_auc_list = []
dir_path = os.path.join(output_root, '%s_checkpoints' % (flag))
if not os.path.exists(dir_path):
os.makedirs(dir_path)
print('==> Preparing data...')
train_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[.5], std=[.5])
])
val_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[.5], std=[.5])
])
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[.5], std=[.5])
])
train_dataset = dataclass[flag](root=input_root, split='train', transform=train_transform, download=download)
train_loader = data.DataLoader(
dataset=train_dataset, batch_size=batch_size, shuffle=True)
val_dataset = dataclass[flag](root=input_root, split='val', transform=val_transform, download=download)
val_loader = data.DataLoader(
dataset=val_dataset, batch_size=batch_size, shuffle=True)
test_dataset = dataclass[flag](root=input_root, split='test', transform=test_transform, download=download)
test_loader = data.DataLoader(
dataset=test_dataset, batch_size=batch_size, shuffle=True)
print('==> Building and training model...')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = ResNet18(in_channels=n_channels, num_classes=n_classes).to(device)
if task == "multi-label, binary-class":
criterion = nn.BCEWithLogitsLoss()
else:
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
for epoch in trange(start_epoch, end_epoch):
train(model, optimizer, criterion, train_loader, device, task)
val(model, val_loader, device, val_auc_list, task, dir_path, epoch)
auc_list = np.array(val_auc_list)
index = auc_list.argmax()
print('epoch %s is the best model' % (index))
print('==> Testing model...')
restore_model_path = os.path.join(dir_path, 'ckpt_%d_auc_%.5f.pth' % (index, auc_list[index]))
model.load_state_dict(torch.load(restore_model_path)['net'])
test(model, 'train', train_loader, device, flag, task, output_root=output_root)
test(model, 'val', val_loader, device, flag, task, output_root=output_root)
test(model, 'test', test_loader, device, flag, task, output_root=output_root)
def train(model, optimizer, criterion, train_loader, device, task):
''' training function
:param model: the model to train
:param optimizer: optimizer used in training
:param criterion: loss function
:param train_loader: DataLoader of training set
:param device: cpu or cuda
:param task: task of current dataset, binary-class/multi-class/multi-label, binary-class
'''
model.train()
for batch_idx, (inputs, targets) in enumerate(train_loader):
optimizer.zero_grad()
outputs = model(inputs.to(device))
if task == 'multi-label, binary-class':
targets = targets.to(torch.float32).to(device)
loss = criterion(outputs, targets)
else:
targets = targets.squeeze().long().to(device)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
def val(model, val_loader, device, val_auc_list, task, dir_path, epoch):
''' validation function
:param model: the model to validate
:param val_loader: DataLoader of validation set
:param device: cpu or cuda
:param val_auc_list: the list to save AUC score of each epoch
:param task: task of current dataset, binary-class/multi-class/multi-label, binary-class
:param dir_path: where to save model
:param epoch: current epoch
'''
model.eval()
y_true = torch.tensor([]).to(device)
y_score = torch.tensor([]).to(device)
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(val_loader):
outputs = model(inputs.to(device))
if task == 'multi-label, binary-class':
targets = targets.to(torch.float32).to(device)
m = nn.Sigmoid()
outputs = m(outputs).to(device)
else:
targets = targets.squeeze().long().to(device)
m = nn.Softmax(dim=1)
outputs = m(outputs).to(device)
targets = targets.float().resize_(len(targets), 1)
y_true = torch.cat((y_true, targets), 0)
y_score = torch.cat((y_score, outputs), 0)
y_true = y_true.cpu().numpy()
y_score = y_score.detach().cpu().numpy()
auc = getAUC(y_true, y_score, task)
val_auc_list.append(auc)
state = {
'net': model.state_dict(),
'auc': auc,
'epoch': epoch,
}
path = os.path.join(dir_path, 'ckpt_%d_auc_%.5f.pth' % (epoch, auc))
torch.save(state, path)
def test(model, split, data_loader, device, flag, task, output_root=None):
''' testing function
:param model: the model to test
:param split: the data to test, 'train/val/test'
:param data_loader: DataLoader of data
:param device: cpu or cuda
:param flag: subset name
:param task: task of current dataset, binary-class/multi-class/multi-label, binary-class
'''
model.eval()
y_true = torch.tensor([]).to(device)
y_score = torch.tensor([]).to(device)
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(data_loader):
outputs = model(inputs.to(device))
if task == 'multi-label, binary-class':
targets = targets.to(torch.float32).to(device)
m = nn.Sigmoid()
outputs = m(outputs).to(device)
else:
targets = targets.squeeze().long().to(device)
m = nn.Softmax(dim=1)
outputs = m(outputs).to(device)
targets = targets.float().resize_(len(targets), 1)
y_true = torch.cat((y_true, targets), 0)
y_score = torch.cat((y_score, outputs), 0)
y_true = y_true.cpu().numpy()
y_score = y_score.detach().cpu().numpy()
auc = getAUC(y_true, y_score, task)
acc = getACC(y_true, y_score, task)
print('%s AUC: %.5f ACC: %.5f' % (split, auc, acc))
if output_root is not None:
output_dir = os.path.join(output_root, flag)
if not os.path.exists(output_dir):
os.mkdir(output_dir)
output_path = os.path.join(output_dir, '%s.csv' % (split))
save_results(y_true, y_score, output_path)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='RUN Baseline model of MedMNIST')
parser.add_argument('--data_name', default='pathmnist', help='subset of MedMNIST', type=str)
parser.add_argument('--input_root', default='./input', help='input root, the source of dataset files', type=str)
parser.add_argument('--output_root', default='./output', help='output root, where to save models and results',
type=str)
parser.add_argument('--num_epoch', default=100, help='num of epochs of training', type=int)
parser.add_argument('--download', default=False, help='whether download the dataset or not', type=bool)
args = parser.parse_args()
data_name = args.data_name.lower()
input_root = args.input_root
output_root = args.output_root
end_epoch = args.num_epoch
download = args.download
main(data_name, input_root, output_root, end_epoch=end_epoch, download=download)
3415
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
