def run(args):
with open(args.model_path + 'cfg.json') as f:
cfg = edict(json.load(f))
device_ids = list(map(int, args.device_ids.split(',')))
num_devices = torch.cuda.device_count()
if num_devices < len(device_ids):
raise Exception('#available gpu : {} < --device_ids : {}'.format(
num_devices, len(device_ids)))
device = torch.device('cuda:{}'.format(device_ids[0]))
ckpt_path = os.path.join(args.model_path, 'best.ckpt')
ckpt = torch.load(ckpt_path, map_location=device)
if args.fl == 'True':
model = Classifier(cfg).to(device).eval()
model.load_state_dict(ckpt['state_dict'])
else:
model = Classifier(cfg)
model = DataParallel(model, device_ids=device_ids).to(device).eval()
model.module.load_state_dict(ckpt['state_dict'])
dataloader_test = DataLoader(ImageDataset(args.in_csv_path,
cfg,
mode='test'),
batch_size=cfg.dev_batch_size,
num_workers=args.num_workers,
drop_last=False,
shuffle=False)
test_epoch(cfg, args, model, dataloader_test, args.out_csv_path)
print('Save best is step :', ckpt['step'], 'AUC :', ckpt['auc_dev_best'])
def build_model(cfg, paramsfile):
model = Classifier(cfg)
model = model.to('cpu')
ckpt = torch.load(paramsfile, map_location='cpu')
state_dict = ckpt['state_dict'] if 'state_dict' in ckpt else ckpt
model.load_state_dict(state_dict)
if 'step' in ckpt and 'auc_dev_best' in ckpt:
print(f"Using model '{paramsfile}' at step: {ckpt['step']} "
f"with AUC: {ckpt['auc_dev_best']}")
return model.eval()
class MDAIModel:
def __init__(self):
root_path = os.path.join(os.path.dirname(os.path.dirname(__file__)),
"model")
with open(os.path.join(root_path, "config/example.json")) as f:
cfg = edict(json.load(f))
self.model = Classifier(cfg)
self.model.cfg.num_classes = [1, 1, 1, 1, 1, 1]
self.model._init_classifier()
self.model._init_attention_map()
self.model._init_bn()
if torch.cuda.is_available():
self.model = self.model.eval().cuda()
else:
self.model = self.model.eval().cpu()
chkpt_path = os.path.join(root_path, "model_best.pt")
self.model.load_state_dict(
torch.load(chkpt_path, map_location=lambda storage, loc: storage))
def predict(self, data):
"""
See https://github.com/mdai/model-deploy/blob/master/mdai/server.py for details on the
schema of `data` and the required schema of the outputs returned by this function.
"""
input_files = data["files"]
input_annotations = data["annotations"]
input_args = data["args"]
outputs = []
for file in input_files:
if file["content_type"] != "application/dicom":
continue
ds = pydicom.dcmread(BytesIO(file["content"]))
x = ds.pixel_array
x_orig = x
# preprocess image
# convert grayscale to RGB
x = cv2.resize(x, (1024, 1024))
x = equalize_adapthist(x.astype(float) / x.max(), clip_limit=0.01)
x = cv2.resize(x, (512, 512))
x = x * 2 - 1
x = np.array([[x, x, x]])
x = torch.from_numpy(x).float()
if torch.cuda.is_available():
x = x.cuda()
else:
x = x.cpu()
with torch.no_grad():
logits, logit_maps = self.model(x)
logits = torch.cat(logits, dim=1).detach().cpu()
y_prob = torch.sigmoid(logits -
torch.from_numpy(threshs).reshape((1,
6)))
y_prob = y_prob.cpu().numpy()
x.requires_grad = True
y_classes = y_prob >= 0.5
class_indices = np.where(y_classes.astype("bool"))[1]
if len(class_indices) == 0:
# no outputs, return 'NONE' output type
output = {
"type": "NONE",
"study_uid": str(ds.StudyInstanceUID),
"series_uid": str(ds.SeriesInstanceUID),
"instance_uid": str(ds.SOPInstanceUID),
"frame_number": None,
}
outputs.append(output)
else:
for class_index in class_indices:
probability = y_prob[0][class_index]
gradcam = GradCam(self.model)
gradcam_output = gradcam.generate_cam(
x, x_orig, class_index)
gradcam_output_buffer = BytesIO()
gradcam_output.save(gradcam_output_buffer, format="PNG")
intgrad = IntegratedGradients(self.model)
intgrad_output = intgrad.generate_integrated_gradients(
x, class_index, 5)
intgrad_output_buffer = BytesIO()
intgrad_output.save(intgrad_output_buffer, format="PNG")
output = {
"type":
"ANNOTATION",
"study_uid":
str(ds.StudyInstanceUID),
"series_uid":
str(ds.SeriesInstanceUID),
"instance_uid":
str(ds.SOPInstanceUID),
"frame_number":
None,
"class_index":
int(class_index),
"data":
None,
"probability":
float(probability),
"explanations": [
{
"name": "Grad-CAM",
"description":
"Visualize how parts of the image affects neural network’s output by looking into the activation maps. From _Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization_ (https://arxiv.org/abs/1610.02391)",
"content": gradcam_output_buffer.getvalue(),
"content_type": "image/png",
},
{
"name": "Integrated Gradients",
"description":
"Visualize an average of the gradients along the construction of the input towards the decision. From _Axiomatic Attribution for Deep Networks_ (https://arxiv.org/abs/1703.01365)",
"content": intgrad_output_buffer.getvalue(),
"content_type": "image/png",
},
],
}
outputs.append(output)
return outputs
def main():
if args.dataset == 'ChestXray-NIHCC':
if args.no_fiding:
classes = [
'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration',
'Mass', 'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation',
'Edema', 'Emphysema', 'Fibrosis', 'Pleural_Thickening',
'Hernia', 'No Fiding'
]
else:
classes = [
'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration',
'Mass', 'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation',
'Edema', 'Emphysema', 'Fibrosis', 'Pleural_Thickening',
'Hernia'
]
elif args.dataset == 'CheXpert-v1.0-small':
classes = [
'No Finding', 'Enlarged Cardiomediastinum', 'Cardiomegaly',
'Lung Opacity', 'Lung Lesion', 'Edema', 'Consolidation',
'Pneumonia', 'Atelectasis', 'Pneumothorax', 'Pleural Effusion',
'Pleural Other', 'Fracture', 'Support Devices'
]
else:
print('--dataset incorrect')
return
torch.manual_seed(args.seed)
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
pin_memory = True if use_gpu else False
print("Initializing dataset: {}".format(args.dataset))
data_transforms = {
'train':
transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.Resize(556),
transforms.CenterCrop(512),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]),
'valid':
transforms.Compose([
transforms.Resize(556),
transforms.CenterCrop(512),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]),
}
datasetTrain = DatasetGenerator(path_base=args.base_dir,
dataset_file='train',
transform=data_transforms['train'],
dataset_=args.dataset,
no_fiding=args.no_fiding)
datasetVal = DatasetGenerator(path_base=args.base_dir,
dataset_file='valid',
transform=data_transforms['valid'],
dataset_=args.dataset,
no_fiding=args.no_fiding)
train_loader = DataLoader(dataset=datasetTrain,
batch_size=args.train_batch,
shuffle=args.train_shuffle,
num_workers=args.workers,
pin_memory=pin_memory)
valid_loader = DataLoader(dataset=datasetVal,
batch_size=args.valid_batch,
shuffle=args.valid_shuffle,
num_workers=args.workers,
pin_memory=pin_memory)
with open(args.infos_densenet) as f:
cfg = edict(json.load(f))
print('Initializing densenet branch')
model_dense = Classifier(cfg)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model_dense.parameters()) / 1000000.0))
with open(args.infos_resnet) as f:
cfg = edict(json.load(f))
print('Initializing resnet branch')
model_res = Classifier(cfg)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model_res.parameters()) / 1000000.0))
print('Initializing fusion branch')
model_fusion = Fusion(input_size=7424, output_size=len(classes))
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model_fusion.parameters()) / 1000000.0))
print("Initializing optimizers")
optimizer_dense = init_optim(args.optim, model_dense.parameters(),
args.learning_rate, args.weight_decay,
args.momentum)
optimizer_res = init_optim(args.optim, model_res.parameters(),
args.learning_rate, args.weight_decay,
args.momentum)
optimizer_fusion = init_optim(args.optim, model_fusion.parameters(),
args.learning_rate, args.weight_decay,
args.momentum)
criterion = nn.BCELoss()
print("Initializing scheduler: {}".format(args.scheduler))
if args.stepsize > 0:
scheduler_dense = init_scheduler(args.scheduler, optimizer_dense,
args.stepsize, args.gamma)
scheduler_res = init_scheduler(args.scheduler, optimizer_res,
args.stepsize, args.gamma)
scheduler_fusion = init_scheduler(args.scheduler, optimizer_fusion,
args.stepsize, args.gamma)
start_epoch = args.start_epoch
best_loss = np.inf
if args.resume_densenet:
checkpoint_dense = torch.load(args.resume_densenet)
model_dense.load_state_dict(checkpoint_dense['state_dict'])
epoch_dense = checkpoint_dense['epoch']
print("Resuming densenet from epoch {}".format(epoch_dense + 1))
if args.resume_resnet:
checkpoint_res = torch.load(args.resume_resnet)
model_res.load_state_dict(checkpoint_res['state_dict'])
epoch_res = checkpoint_res['epoch']
print("Resuming resnet from epoch {}".format(epoch_res + 1))
if args.resume_fusion:
checkpoint_fusion = torch.load(args.resume_fusion)
model_fusion.load_state_dict(checkpoint_fusion['state_dict'])
epoch_fusion = checkpoint_fusion['epoch']
print("Resuming fusion from epoch {}".format(epoch_fusion + 1))
if use_gpu:
model_dense = nn.DataParallel(model_dense).cuda()
model_res = nn.DataParallel(model_res).cuda()
model_fusion = nn.DataParallel(model_fusion).cuda()
if args.evaluate:
print("Evaluate only")
if args.step == 1:
valid('step1', model_dense, model_res, model_fusion, valid_loader,
criterion, args.print_freq, classes, cfg,
data_transforms['valid'])
elif args.step == 2:
valid('step2', model_dense, model_res, model_fusion, valid_loader,
criterion, args.print_freq, classes, cfg,
data_transforms['valid'])
elif args.step == 3:
valid('step3', model_dense, model_res, model_fusion, valid_loader,
criterion, args.print_freq, classes, cfg,
data_transforms['valid'])
else:
print('args.step not found')
return
if args.step == 1:
#################################### DENSENET BRANCH INIT ##########################################
start_time = time.time()
train_time = 0
best_epoch = 0
print("==> Start training of densenet branch")
for p in model_dense.parameters():
p.requires_grad = True
for p in model_res.parameters():
p.requires_grad = False
for p in model_fusion.parameters():
p.requires_grad = True
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train('step1', model_dense, model_res, model_fusion, train_loader,
optimizer_dense, optimizer_res, optimizer_fusion, criterion,
args.print_freq, epoch, args.max_epoch, cfg,
data_transforms['train'])
train_time += round(time.time() - start_train_time)
if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (
epoch + 1) == args.max_epoch:
print("==> Validation")
loss_val = valid('step1', model_dense, model_res, model_fusion,
valid_loader, criterion, args.print_freq,
classes, cfg, data_transforms['valid'])
if args.stepsize > 0:
if args.scheduler == 'ReduceLROnPlateau':
scheduler_dense.step(loss_val)
scheduler_fusion.step(loss_val)
else:
scheduler_dense.step()
scheduler_fusion.step()
is_best = loss_val < best_loss
if is_best:
best_loss = loss_val
best_epoch = epoch + 1
if use_gpu:
state_dict_dense = model_dense.module.state_dict()
state_dict_fusion = model_fusion.module.state_dict()
else:
state_dict_dense = model_dense.state_dict()
state_dict_fusion = model_fusion.state_dict()
save_checkpoint(
{
'state_dict': state_dict_dense,
'loss': best_loss,
'epoch': epoch,
}, is_best, args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'dense')
save_checkpoint(
{
'state_dict': state_dict_fusion,
'loss': best_loss,
'epoch': epoch,
}, is_best, args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'fusion')
print("==> Best Validation Loss {:.4%}, achieved at epoch {}".format(
best_loss, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Dense branch finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}."
.format(elapsed, train_time))
#################################### DENSENET BRANCH END ##########################################
elif args.step == 2:
#################################### RESNET BRANCH INIT ##########################################
start_time = time.time()
train_time = 0
best_epoch = 0
print("==> Start training of local branch")
for p in model_dense.parameters():
p.requires_grad = False
for p in model_res.parameters():
p.requires_grad = True
for p in model_fusion.parameters():
p.requires_grad = True
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train('step2', model_dense, model_res, model_fusion, train_loader,
optimizer_dense, optimizer_res, optimizer_fusion, criterion,
args.print_freq, epoch, args.max_epoch, cfg,
data_transforms['train'])
train_time += round(time.time() - start_train_time)
if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (
epoch + 1) == args.max_epoch:
print("==> Validation")
loss_val = valid('step2', model_dense, model_res, model_fusion,
valid_loader, criterion, args.print_freq,
classes, cfg, data_transforms['valid'])
if args.stepsize > 0:
if args.scheduler == 'ReduceLROnPlateau':
scheduler_res.step(loss_val)
scheduler_fusion.step(loss_val)
else:
scheduler_res.step()
scheduler_fusion.step()
is_best = loss_val < best_loss
if is_best:
best_loss = loss_val
best_epoch = epoch + 1
if use_gpu:
state_dict_res = model_res.module.state_dict()
state_dict_fusion = model_fusion.module.state_dict()
else:
state_dict_res = model_res.state_dict()
state_dict_fusion = model_fusion.state_dict()
save_checkpoint(
{
'state_dict': state_dict_res,
'loss': best_loss,
'epoch': epoch,
}, is_best, args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'res')
save_checkpoint(
{
'state_dict': state_dict_fusion,
'loss': best_loss,
'epoch': epoch,
}, is_best, args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'fusion')
print("==> Best Validation Loss {:.4%}, achieved at epoch {}".format(
best_loss, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Resnet branch finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}."
.format(elapsed, train_time))
#################################### RESNET BRANCH END ##########################################
elif args.step == 3:
#################################### FUSION BRANCH INIT ##########################################
start_time = time.time()
train_time = 0
best_epoch = 0
print("==> Start training of fusion branch")
for p in model_dense.parameters():
p.requires_grad = True
for p in model_res.parameters():
p.requires_grad = True
for p in model_fusion.parameters():
p.requires_grad = True
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train('step3', model_dense, model_res, model_fusion, train_loader,
optimizer_dense, optimizer_res, optimizer_fusion, criterion,
args.print_freq, epoch, args.max_epoch, cfg,
data_transforms['train'])
train_time += round(time.time() - start_train_time)
if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (
epoch + 1) == args.max_epoch:
print("==> Validation")
loss_val = valid('step3', model_dense, model_res, model_fusion,
valid_loader, criterion, args.print_freq,
classes, cfg, data_transforms['valid'])
if args.stepsize > 0:
if args.scheduler == 'ReduceLROnPlateau':
scheduler_dense.step(loss_val)
scheduler_res.step(loss_val)
scheduler_fusion.step(loss_val)
else:
scheduler_dense.step()
scheduler_res.step()
scheduler_fusion.step()
is_best = loss_val < best_loss
if is_best:
best_loss = loss_val
best_epoch = epoch + 1
if use_gpu:
state_dict_dense = model_dense.module.state_dict()
state_dict_res = model_res.module.state_dict()
state_dict_fusion = model_fusion.module.state_dict()
else:
state_dict_dense = model_dense.state_dict()
state_dict_res = model_res.state_dict()
state_dict_fusion = model_fusion.state_dict()
save_checkpoint(
{
'state_dict': state_dict_dense,
'loss': best_loss,
'epoch': epoch,
}, is_best, args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'dense')
save_checkpoint(
{
'state_dict': state_dict_res,
'loss': best_loss,
'epoch': epoch,
}, is_best, args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'res')
save_checkpoint(
{
'state_dict': state_dict_fusion,
'loss': best_loss,
'epoch': epoch,
}, is_best, args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'fusion')
print("==> Best Validation Loss {:.4%}, achieved at epoch {}".format(
best_loss, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Fusion branch finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}."
.format(elapsed, train_time))
#################################### FUSION BRANCH END ##########################################
else:
print('args.step not found')
def run_fl(args):
with open(args.cfg_path) as f:
cfg = edict(json.load(f))
if args.verbose is True:
print(json.dumps(cfg, indent=4))
if not os.path.exists(args.save_path):
os.mkdir(args.save_path)
if args.logtofile is True:
logging.basicConfig(filename=args.save_path + '/log.txt',
filemode="w",
level=logging.INFO)
else:
logging.basicConfig(level=logging.INFO)
if not args.resume:
with open(os.path.join(args.save_path, 'cfg.json'), 'w') as f:
json.dump(cfg, f, indent=1)
device_ids = list(map(int, args.device_ids.split(',')))
num_devices = torch.cuda.device_count()
if num_devices < len(device_ids):
raise Exception('#available gpu : {} < --device_ids : {}'.format(
num_devices, len(device_ids)))
device = torch.device('cuda:{}'.format(device_ids[0]))
# initialise global model
model = Classifier(cfg).to(device).train()
if args.verbose is True:
from torchsummary import summary
if cfg.fix_ratio:
h, w = cfg.long_side, cfg.long_side
else:
h, w = cfg.height, cfg.width
summary(model.to(device), (3, h, w))
if args.pre_train is not None:
if os.path.exists(args.pre_train):
ckpt = torch.load(args.pre_train, map_location=device)
model.load_state_dict(ckpt)
src_folder = os.path.dirname(os.path.abspath(__file__)) + '/../'
dst_folder = os.path.join(args.save_path, 'classification')
rc, size = subprocess.getstatusoutput('du --max-depth=0 %s | cut -f1' %
src_folder)
if rc != 0:
raise Exception('Copy folder error : {}'.format(rc))
else:
print('Successfully determined size of directory')
rc, err_msg = subprocess.getstatusoutput('cp -R %s %s' %
(src_folder, dst_folder))
if rc != 0:
raise Exception('copy folder error : {}'.format(err_msg))
else:
print('Successfully copied folder')
# copy train files
train_files = cfg.train_csv
clients = {}
for i, c in enumerate(string.ascii_uppercase):
if i < len(train_files):
clients[c] = {}
else:
break
# initialise clients
for i, client in enumerate(clients):
copyfile(train_files[i],
os.path.join(args.save_path, f'train_{client}.csv'))
clients[client]['dataloader_train'] =\
DataLoader(
ImageDataset(train_files[i], cfg, mode='train'),
batch_size=cfg.train_batch_size,
num_workers=args.num_workers,drop_last=True,
shuffle=True
)
clients[client]['bytes_uploaded'] = 0.0
clients[client]['epoch'] = 0
copyfile(cfg.dev_csv, os.path.join(args.save_path, 'dev.csv'))
dataloader_dev = DataLoader(ImageDataset(cfg.dev_csv, cfg, mode='dev'),
batch_size=cfg.dev_batch_size,
num_workers=args.num_workers,
drop_last=False,
shuffle=False)
dev_header = dataloader_dev.dataset._label_header
w_global = model.state_dict()
summary_train = {'epoch': 0, 'step': 0}
summary_dev = {'loss': float('inf'), 'acc': 0.0}
summary_writer = SummaryWriter(args.save_path)
comm_rounds = cfg.epoch
best_dict = {
"acc_dev_best": 0.0,
"auc_dev_best": 0.0,
"loss_dev_best": float('inf'),
"fused_dev_best": 0.0,
"best_idx": 1
}
# Communication rounds loop
for cr in range(comm_rounds):
logging.info('{}, Start communication round {} of FL - {} ...'.format(
time.strftime("%Y-%m-%d %H:%M:%S"), cr + 1, cfg.fl_technique))
w_locals = []
for client in clients:
logging.info(
'{}, Start local training process for client {}, communication round: {} ...'
.format(time.strftime("%Y-%m-%d %H:%M:%S"), client, cr + 1))
# Load previous current global model as start point
model = Classifier(cfg).to(device).train()
model.load_state_dict(w_global)
if cfg.fl_technique == "FedProx":
global_weight_collector = get_global_weights(model, device)
else:
global_weight_collector = None
optimizer = get_optimizer(model.parameters(), cfg)
# local training loops
for epoch in range(cfg.local_epoch):
lr = lr_schedule(cfg.lr, cfg.lr_factor, epoch, cfg.lr_epochs)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
summary_train, best_dict = train_epoch_fl(
summary_train, summary_dev, cfg, args, model,
clients[client]['dataloader_train'], dataloader_dev,
optimizer, summary_writer, best_dict, dev_header, epoch,
global_weight_collector)
summary_train['step'] += 1
bytes_to_upload = sys.getsizeof(model.state_dict())
clients[client]['bytes_uploaded'] += bytes_to_upload
logging.info(
'{}, Completed local rounds for client {} in communication round {}. '
'Uploading {} bytes to server, {} bytes in total sent from client'
.format(time.strftime("%Y-%m-%d %H:%M:%S"), client, cr + 1,
bytes_to_upload, clients[client]['bytes_uploaded']))
w_locals.append(model.state_dict())
if cfg.fl_technique == "FedAvg":
w_global = fed_avg(w_locals)
elif cfg.fl_technique == 'WFedAvg':
w_global = weighted_fed_avg(w_locals, cfg.train_proportions)
elif cfg.fl_technique == 'FedProx':
# Use weighted FedAvg when using FedProx
w_global = weighted_fed_avg(w_locals, cfg.train_proportions)
# Test the performance of the averaged model
avged_model = Classifier(cfg).to(device)
avged_model.load_state_dict(w_global)
time_now = time.time()
summary_dev, predlist, true_list = test_epoch(summary_dev, cfg, args,
avged_model,
dataloader_dev)
time_spent = time.time() - time_now
auclist = []
for i in range(len(cfg.num_classes)):
y_pred = predlist[i]
y_true = true_list[i]
fpr, tpr, thresholds = metrics.roc_curve(y_true,
y_pred,
pos_label=1)
auc = metrics.auc(fpr, tpr)
auclist.append(auc)
auc_summary = np.array(auclist)
loss_dev_str = ' '.join(
map(lambda x: '{:.5f}'.format(x), summary_dev['loss']))
acc_dev_str = ' '.join(
map(lambda x: '{:.3f}'.format(x), summary_dev['acc']))
auc_dev_str = ' '.join(map(lambda x: '{:.3f}'.format(x), auc_summary))
logging.info(
'{}, Averaged Model -> Dev, Step : {}, Loss : {}, Acc : {}, Auc : {},'
'Mean auc: {:.3f} '
'Run Time : {:.2f} sec'.format(time.strftime("%Y-%m-%d %H:%M:%S"),
summary_train['step'], loss_dev_str,
acc_dev_str, auc_dev_str,
auc_summary.mean(), time_spent))
for t in range(len(cfg.num_classes)):
summary_writer.add_scalar('dev/loss_{}'.format(dev_header[t]),
summary_dev['loss'][t],
summary_train['step'])
summary_writer.add_scalar('dev/acc_{}'.format(dev_header[t]),
summary_dev['acc'][t],
summary_train['step'])
summary_writer.add_scalar('dev/auc_{}'.format(dev_header[t]),
auc_summary[t], summary_train['step'])
save_best = False
mean_acc = summary_dev['acc'][cfg.save_index].mean()
if mean_acc >= best_dict['acc_dev_best']:
best_dict['acc_dev_best'] = mean_acc
if cfg.best_target == 'acc':
save_best = True
mean_auc = auc_summary[cfg.save_index].mean()
if mean_auc >= best_dict['auc_dev_best']:
best_dict['auc_dev_best'] = mean_auc
if cfg.best_target == 'auc':
save_best = True
mean_loss = summary_dev['loss'][cfg.save_index].mean()
if mean_loss <= best_dict['loss_dev_best']:
best_dict['loss_dev_best'] = mean_loss
if cfg.best_target == 'loss':
save_best = True
if save_best:
torch.save(
{
'epoch': summary_train['epoch'],
'step': summary_train['step'],
'acc_dev_best': best_dict['acc_dev_best'],
'auc_dev_best': best_dict['auc_dev_best'],
'loss_dev_best': best_dict['loss_dev_best'],
'state_dict': avged_model.state_dict()
},
os.path.join(args.save_path,
'best{}.ckpt'.format(best_dict['best_idx'])))
best_dict['best_idx'] += 1
if best_dict['best_idx'] > cfg.save_top_k:
best_dict['best_idx'] = 1
logging.info('{}, Best, Step : {}, Loss : {}, Acc : {},'
'Auc :{},Best Auc : {:.3f}'.format(
time.strftime("%Y-%m-%d %H:%M:%S"),
summary_train['step'], loss_dev_str, acc_dev_str,
auc_dev_str, best_dict['auc_dev_best']))
torch.save(
{
'epoch': cr,
'step': summary_train['step'],
'acc_dev_best': best_dict['acc_dev_best'],
'auc_dev_best': best_dict['auc_dev_best'],
'loss_dev_best': best_dict['loss_dev_best'],
'state_dict': avged_model.state_dict()
}, os.path.join(args.save_path, 'train.ckpt'))
class MDAIModel:
def __init__(self):
root_path = os.path.dirname(__file__)
with open(os.path.join(root_path, "config/example.json")) as f:
cfg = edict(json.load(f))
self.model = Classifier(cfg)
self.model.cfg.num_classes = [1, 1, 1, 1, 1, 1]
self.model._init_classifier()
self.model._init_attention_map()
self.model._init_bn()
if torch.cuda.is_available():
self.model = self.model.eval().cuda()
else:
self.model = self.model.eval().cpu()
chkpt_path = os.path.join(root_path, "model_best.pt")
self.model.load_state_dict(
torch.load(chkpt_path, map_location=lambda storage, loc: storage)
)
def predict(self, data):
"""
The input data has the following schema:
{
"instances": [
{
"file": "bytes"
"tags": {
"StudyInstanceUID": "str",
"SeriesInstanceUID": "str",
"SOPInstanceUID": "str",
...
}
},
...
],
"args": {
"arg1": "str",
"arg2": "str",
...
}
}
Model scope specifies whether an entire study, series, or instance is given to the model.
If the model scope is 'INSTANCE', then `instances` will be a single instance (list length of 1).
If the model scope is 'SERIES', then `instances` will be a list of all instances in a series.
If the model scope is 'STUDY', then `instances` will be a list of all instances in a study.
The additional `args` dict supply values that may be used in a given run.
For a single instance dict, `files` is the raw binary data representing a DICOM file, and
can be loaded using: `ds = pydicom.dcmread(BytesIO(instance["file"]))`.
The results returned by this function should have the following schema:
[
{
"type": "str", // 'NONE', 'ANNOTATION', 'IMAGE', 'DICOM', 'TEXT'
"study_uid": "str",
"series_uid": "str",
"instance_uid": "str",
"frame_number": "int",
"class_index": "int",
"data": {},
"probability": "float",
"explanations": [
{
"name": "str",
"description": "str",
"content": "bytes",
"content_type": "str",
},
...
],
},
...
]
The DICOM UIDs must be supplied based on the scope of the label attached to `class_index`.
"""
input_instances = data["instances"]
input_args = data["args"]
results = []
for instance in input_instances:
tags = instance["tags"]
ds = pydicom.dcmread(BytesIO(instance["file"]))
x = ds.pixel_array
x_orig = x
# preprocess image
# convert grayscale to RGB
x = cv2.resize(x, (1024, 1024))
x = equalize_adapthist(x.astype(float) / x.max(), clip_limit=0.01)
x = cv2.resize(x, (512, 512))
x = x * 2 - 1
x = np.array([[x, x, x]])
x = torch.from_numpy(x).float()
if torch.cuda.is_available():
x = x.cuda()
else:
x = x.cpu()
with torch.no_grad():
logits, logit_maps = self.model(x)
logits = torch.cat(logits, dim=1).detach().cpu()
y_prob = torch.sigmoid(logits - torch.from_numpy(threshs).reshape((1, 6)))
y_prob = y_prob.cpu().numpy()
x.requires_grad = True
y_classes = y_prob >= 0.5
class_indices = np.where(y_classes.astype("bool"))[1]
if len(class_indices) == 0:
# no outputs, return 'NONE' output type
result = {
"type": "NONE",
"study_uid": tags["StudyInstanceUID"],
"series_uid": tags["SeriesInstanceUID"],
"instance_uid": tags["SOPInstanceUID"],
"frame_number": None,
}
results.append(result)
else:
for class_index in class_indices:
probability = y_prob[0][class_index]
gradcam = GradCam(self.model)
gradcam_output = gradcam.generate_cam(x, x_orig, class_index)
gradcam_output_buffer = BytesIO()
gradcam_output.save(gradcam_output_buffer, format="PNG")
intgrad = IntegratedGradients(self.model)
intgrad_output = intgrad.generate_integrated_gradients(x, class_index, 5)
intgrad_output_buffer = BytesIO()
intgrad_output.save(intgrad_output_buffer, format="PNG")
result = {
"type": "ANNOTATION",
"study_uid": tags["StudyInstanceUID"],
"series_uid": tags["SeriesInstanceUID"],
"instance_uid": tags["SOPInstanceUID"],
"frame_number": None,
"class_index": int(class_index),
"data": None,
"probability": float(probability),
"explanations": [
{
"name": "Grad-CAM",
"description": "Visualize how parts of the image affects neural network’s output by looking into the activation maps. From _Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization_ (https://arxiv.org/abs/1610.02391)",
"content": gradcam_output_buffer.getvalue(),
"content_type": "image/png",
},
{
"name": "Integrated Gradients",
"description": "Visualize an average of the gradients along the construction of the input towards the decision. From _Axiomatic Attribution for Deep Networks_ (https://arxiv.org/abs/1703.01365)",
"content": intgrad_output_buffer.getvalue(),
"content_type": "image/png",
},
],
}
results.append(result)
return results
from converter.pytorch.pytorch_parser import PytorchParser # noqa
from model.classifier import Classifier # noqa
parser = argparse.ArgumentParser(description='test converter')
parser.add_argument('model_path', default=None, metavar='MODEL_PATH', type=str,
help="Path to the trained models")
args = parser.parse_args()
with open(args.model_path+'cfg.json') as f:
cfg = edict(json.load(f))
model_file = "model/best.pth"
device = torch.device('cpu') # PyTorch v0.4.0
net = Classifier(cfg)
ckpt = torch.load("model/best.ckpt")
net.load_state_dict(ckpt['state_dict'], strict=False)
torch.save(net, model_file)
net.eval()
dummy_input = torch.ones([1, 3, 1024, 1024])
net.to(device)
output = net(dummy_input)
device = torch.device("cuda") # PyTorch v0.4.0
summary(net.to(device), (3, 1024, 1024))
pytorch_parser = PytorchParser(model_file, [3, 1024, 1024])
#
pytorch_parser.run(model_file)
