2021-1-capstone-design1 / KST_Project1

Go to a project
Toggle navigation Toggle navigation pinning
KimJyun
d94d4a55 1 parent 48e03165

최종 코드 업로드

Inline Side-by-side
Showing 15 changed files with 1190 additions and 0 deletions
This diff could not be displayed because it is too large.
import pandas as pd
import torch
import numpy as np
from torch.utils.data import Dataset
import os
from PIL import Image
class CXRDataset(Dataset):
def __init__(
self,
path_to_images,
fold,
transform=None,
transform_bb=None,
finding="any",
fine_tune=False,
regression=False,
label_path="/content/gdrive/MyDrive/ColabNotebooks/brixia/labels"):
self.transform = transform
self.transform_bb = transform_bb
self.path_to_images = path_to_images
if not fine_tune:
self.df = pd.read_csv(label_path + "/nih_original_split.csv")
elif fine_tune and not regression:
self.df = pd.read_csv(label_path + "/brixia_split_classification.csv")
else:
self.df = pd.read_csv(label_path + "/brixia_split_regression.csv")
self.fold = fold
self.fine_tune = fine_tune
self.regression = regression
if not fold == 'BBox':
self.df = self.df[self.df['fold'] == fold]
else:
bbox_images_df = pd.read_csv(label_path + "/BBox_List_2017.csv")
self.df = pd.merge(left=self.df, right=bbox_images_df, how="inner", on="Image Index")
if not self.fine_tune:
self.PRED_LABEL = [
'Atelectasis',
'Cardiomegaly',
'Effusion',
'Infiltration',
'Mass',
'Nodule',
'Pneumonia',
'Pneumothorax',
'Consolidation',
'Edema',
'Emphysema',
'Fibrosis',
'Pleural_Thickening',
'Hernia']
else:
self.PRED_LABEL = [
'Detector01',
'Detector2',
'Detector3']
if not finding == "any" and not fine_tune: # can filter for positive findings of the kind described; useful for evaluation
self.df = self.df[self.df['Finding Label'] == finding]
elif not finding == "any" and fine_tune and not regression:
self.df = self.df[self.df[finding] == 1]
self.df = self.df.set_index("Image Index")
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
image = Image.open(
os.path.join(
self.path_to_images,
self.df.index[idx]))
image = image.convert('RGB')
if not self.fine_tune:
label = np.zeros(len(self.PRED_LABEL), dtype=int)
for i in range(0, len(self.PRED_LABEL)):
# can leave zero if zero, else make one
if self.df[self.PRED_LABEL[i].strip()].iloc[idx].astype('int') > 0:
label[i] = self.df[self.PRED_LABEL[i].strip()
].iloc[idx].astype('int')
elif self.fine_tune and not self.regression:
covid_label = np.zeros(len(self.PRED_LABEL), dtype=int)
covid_label[0] = self.df['Detector01'].iloc[idx]
covid_label[1] = self.df['Detector2'].iloc[idx]
covid_label[2] = self.df['Detector3'].iloc[idx]
else:
ground_truth = np.array(self.df['BrixiaScoreGlobal'].iloc[idx].astype('float32'))
if self.transform:
image = self.transform(image)
if self.fold == "BBox":
# exctract bounding box coordinates from dataframe, they exist in the the columns specified below
bounding_box = self.df.iloc[idx, -7:-3].to_numpy()
if self.transform_bb:
transformed_bounding_box = self.transform_bb(bounding_box)
return image, label, self.df.index[idx], transformed_bounding_box
elif self.fine_tune and not self.regression:
return image, covid_label, self.df.index[idx]
elif self.fine_tune and self.regression:
return image, ground_truth, self.df.index[idx]
else:
return image, label, self.df.index[idx]
def pos_neg_balance_weights(self):
pos_neg_weights = []
for i in range(0, len(self.PRED_LABEL)):
num_negatives = self.df[self.df[self.PRED_LABEL[i].strip()] == 0].shape[0]
num_positives = self.df[self.df[self.PRED_LABEL[i].strip()] == 1].shape[0]
pos_neg_weights.append(num_negatives / num_positives)
pos_neg_weights = torch.Tensor(pos_neg_weights)
pos_neg_weights = pos_neg_weights.cuda()
pos_neg_weights = pos_neg_weights.type(torch.cuda.FloatTensor)
return pos_neg_weights
class RescaleBB(object):
"""Rescale the bounding box in a sample to a given size.
Args:
output_image_size (int): Desired output size.
"""
def __init__(self, output_image_size, original_image_size):
assert isinstance(output_image_size, int)
self.output_image_size = output_image_size
self.original_image_size = original_image_size
def __call__(self, sample):
assert sample.shape == (4,)
x, y, w, h = sample[0], sample[1], sample[2], sample[3]
scale_factor = self.output_image_size / self.original_image_size
new_x, new_y, new_w, new_h = x * scale_factor, y * scale_factor, w * scale_factor, h * scale_factor
transformed_sample = np.array([new_x, new_y, new_w, new_h])
return transformed_sample
class BrixiaScoreLocal:
def __init__(self, label_path):
self.data_brixia = pd.read_csv(label_path + "/metadata_global_v2.csv", sep=";")
self.data_brixia.set_index("Filename", inplace=True)
def getScore(self, filename,print_score=False):
score = self.data_brixia.loc[filename.replace(".jpg", ".dcm"), "BrixiaScore"].astype(str)
score = '0' * (6 - len(score)) + score
if print_score:
print('Brixia 6 regions Score: ')
print(score[0], ' | ', score[3])
print(score[1], ' | ', score[4])
print(score[2], ' | ', score[5])
return list(map(int, score))
import torch
import pandas as pd
import cxr_dataset as CXR
from torch.utils.data import Dataset, DataLoader
import sklearn.metrics as sklm
import numpy as np
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def make_pred_multilabel(dataloader, model, save_as_csv=False, fine_tune=False):
"""
Gives predictions for test fold and calculates AUCs using previously trained model
Args:
data_transforms: torchvision transforms to preprocess raw images; same as validation transforms
model: densenet-121 from torchvision previously fine tuned to training data
PATH_TO_IMAGES: path at which NIH images can be found
Returns:
pred_df: dataframe containing individual predictions and ground truth for each test image
auc_df: dataframe containing aggregate AUCs by train/test tuples
"""
batch_size = dataloader.batch_size
# set model to eval mode; required for proper predictions given use of batchnorm
model.train(False)
# create empty dfs
pred_df = pd.DataFrame(columns=["Image Index"])
true_df = pd.DataFrame(columns=["Image Index"])
# iterate over dataloader
for i, data in enumerate(dataloader):
inputs, labels, _ = data
inputs, labels = inputs.to(device), labels.to(device)
true_labels = labels.cpu().data.numpy()
# batch_size = true_labels.shape
outputs = model(inputs)
outputs = torch.sigmoid(outputs)
probs = outputs.cpu().data.numpy()
# get predictions and true values for each item in batch
for j in range(0, true_labels.shape[0]):
thisrow = {}
truerow = {}
thisrow["Image Index"] = dataloader.dataset.df.index[batch_size * i + j]
truerow["Image Index"] = dataloader.dataset.df.index[batch_size * i + j]
# iterate over each entry in prediction vector; each corresponds to
# individual label
for k in range(len(dataloader.dataset.PRED_LABEL)):
thisrow["prob_" + dataloader.dataset.PRED_LABEL[k]] = probs[j, k]
truerow[dataloader.dataset.PRED_LABEL[k]] = true_labels[j, k]
pred_df = pred_df.append(thisrow, ignore_index=True)
true_df = true_df.append(truerow, ignore_index=True)
# if(i % 10 == 0):
# print(str(i * BATCH_SIZE))
auc_df = pd.DataFrame(columns=["label", "auc"])
# calc AUCs
for column in true_df:
if not fine_tune:
if column not in [
'Atelectasis',
'Cardiomegaly',
'Effusion',
'Infiltration',
'Mass',
'Nodule',
'Pneumonia',
'Pneumothorax',
'Consolidation',
'Edema',
'Emphysema',
'Fibrosis',
'Pleural_Thickening',
'Hernia']:
continue
else:
if column not in [
'Detector01',
'Detector2',
'Detector3']:
continue
actual = true_df[column]
pred = pred_df["prob_" + column]
thisrow = {}
thisrow['label'] = column
thisrow['auc'] = np.nan
thisrow['AP'] = np.nan
try:
thisrow['auc'] = sklm.roc_auc_score(actual.to_numpy().astype(int), pred.to_numpy())
thisrow['AP'] = sklm.average_precision_score(actual.to_numpy().astype(int), pred.to_numpy())
except BaseException:
print("can't calculate auc for " + str(column))
auc_df = auc_df.append(thisrow, ignore_index=True)
if save_as_csv:
pred_df.to_csv("results/preds.csv", index=False)
auc_df.to_csv("results/aucs.csv", index=False)
return pred_df, auc_df
def evaluate_mae(dataloader, model):
"""
Calculates MAE using previously trained model
Args:
data_transforms: torchvision transforms to preprocess raw images; same as validation transforms
model: densenet-121 from torchvision previously fine tuned to training data
Returns:
mae: MAE
"""
# calc preds in batches of 32, can reduce if your GPU has less RAM
batch_size = dataloader.batch_size
# set model to eval mode; required for proper predictions given use of batchnorm
model.train(False)
# create empty dfs
pred_df = pd.DataFrame(columns=["Image Index"])
true_df = pd.DataFrame(columns=["Image Index"])
# iterate over dataloader
for i, data in enumerate(dataloader):
inputs, ground_truths, _ = data
inputs, ground_truths = inputs.to(device), ground_truths.to(device)
true_scores = ground_truths.cpu().data.numpy()
outputs = model(inputs)
preds = outputs.cpu().data.numpy()
# get predictions and true values for each item in batch
for j in range(0, true_scores.shape[0]):
thisrow = {}
truerow = {}
thisrow["Image Index"] = dataloader.dataset.df.index[batch_size * i + j]
truerow["Image Index"] = dataloader.dataset.df.index[batch_size * i + j]
# iterate over each entry in prediction vector; each corresponds to
# individual label
thisrow["pred_score"] = preds[j]
truerow["true_score"] = true_scores[j]
pred_df = pred_df.append(thisrow, ignore_index=True)
true_df = true_df.append(truerow, ignore_index=True)
actual = true_df["true_score"]
pred = pred_df["pred_score"]
try:
mae = sklm.mean_absolute_error(actual.to_numpy().astype(int), pred.to_numpy())
return mae, true_df, pred_df
except BaseException:
print("can't calculate mae")
from __future__ import print_function, division
# pytorch imports
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torchvision import transforms, utils
from tensorboardX import SummaryWriter
# general imports
import os
import time
from shutil import rmtree
# data science imports
import csv
import cxr_dataset as CXR
import eval_model as E
use_gpu = torch.cuda.is_available()
gpu_count = torch.cuda.device_count()
print("Available GPU count:" + str(gpu_count))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def checkpoint(model, best_loss, epoch, LR, filename):
"""
Saves checkpoint of torchvision model during training.
Args:
model: torchvision model to be saved
best_loss: best val loss achieved so far in training
epoch: current epoch of training
LR: current learning rate in training
Returns:
None
"""
print('saving')
state = {
'model': model,
'best_loss': best_loss,
'epoch': epoch,
'rng_state': torch.get_rng_state(),
'LR': LR
}
torch.save(state, 'results/' + filename)
def pos_neg_weights_in_batch(labels_batch):
num_total = labels_batch.shape[0] * labels_batch.shape[1]
num_positives = labels_batch.sum()
num_negatives = num_total - num_positives
if not num_positives == 0:
beta_p = num_negatives / num_positives
else:
beta_p = num_negatives
beta_p = torch.as_tensor(beta_p)
beta_p = beta_p.to(device)
beta_p = beta_p.type(torch.cuda.FloatTensor)
return beta_p
def train_model(
model,
criterion,
optimizer,
LR,
num_epochs,
dataloaders,
dataset_sizes,
weight_decay,
weighted_cross_entropy_batchwise=False,
fine_tune=False,
regression=False):
"""
Fine tunes torchvision model to NIH CXR data.
Args:
model: torchvision model to be finetuned (densenet-121 in this case)
criterion: loss criterion (binary cross entropy loss, BCELoss)
optimizer: optimizer to use in training (SGD)
LR: learning rate
num_epochs: continue training up to this many epochs
dataloaders: pytorch train and val dataloaders
dataset_sizes: length of train and val datasets
weight_decay: weight decay parameter we use in SGD with momentum
Returns:
model: trained torchvision model
best_epoch: epoch on which best model val loss was obtained
"""
since = time.time()
start_epoch = 1
best_loss = 999999
best_epoch = -1
last_train_loss = -1
tensorboard_writer_train = SummaryWriter('runs/loss/train_loss')
tensorboard_writer_val = SummaryWriter('runs/loss/val_loss')
if not fine_tune:
PRED_LABEL = [
'Atelectasis',
'Cardiomegaly',
'Effusion',
'Infiltration',
'Mass',
'Nodule',
'Pneumonia',
'Pneumothorax',
'Consolidation',
'Edema',
'Emphysema',
'Fibrosis',
'Pleural_Thickening',
'Hernia']
else:
PRED_LABEL = [
'Detector01',
'Detector2',
'Detector3']
if not regression:
tensorboard_writer_auc = {}
tensorboard_writer_AP = {}
for label in PRED_LABEL:
tensorboard_writer_auc[label] = SummaryWriter('runs/auc/'+label)
tensorboard_writer_AP[label] = SummaryWriter('runs/ap/' + label)
else:
tensorboard_writer_mae = SummaryWriter('runs/mae')
# iterate over epochs
for epoch in range(start_epoch, num_epochs + 1):
print('Epoch {}/{}'.format(epoch, num_epochs))
print('-' * 10)
# set model to train or eval mode based on whether we are in train or
# val; necessary to get correct predictions given batchnorm
for phase in ['train', 'val']:
if phase == 'train':
model.train(True)
else:
model.train(False)
running_loss = 0.0
total_done = 0
for data in dataloaders[phase]:
if not regression:
inputs, labels, _ = data
else:
inputs, ground_truths, _ = data
batch_size = inputs.shape[0]
inputs = inputs.to(device)
if not regression:
labels = (labels.to(device)).float()
else:
ground_truths = (ground_truths.to(device)).float()
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
# calculate gradient and update parameters in train phase
optimizer.zero_grad()
if weighted_cross_entropy_batchwise:
beta = pos_neg_weights_in_batch(labels)
criterion = nn.BCEWithLogitsLoss(pos_weight=beta)
if not regression:
loss = criterion(outputs, labels)
else:
ground_truths = ground_truths.unsqueeze(1)
loss = criterion(outputs, ground_truths)
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.item() * batch_size
epoch_loss = running_loss / dataset_sizes[phase]
if phase == 'train':
tensorboard_writer_train.add_scalar('Loss', epoch_loss, epoch)
last_train_loss = epoch_loss
elif phase == 'val':
tensorboard_writer_val.add_scalar('Loss', epoch_loss, epoch)
if not regression:
preds, aucs = E.make_pred_multilabel(dataloaders['val'], model, save_as_csv=False, fine_tune=fine_tune)
aucs.set_index('label', inplace=True)
print(aucs)
for label in PRED_LABEL:
tensorboard_writer_auc[label].add_scalar('AUC', aucs.loc[label, 'auc'], epoch)
tensorboard_writer_AP[label].add_scalar('AP', aucs.loc[label, 'AP'], epoch)
else:
mae, _, _ = E.evaluate_mae(dataloaders['val'], model)
print('MAE: ', mae)
tensorboard_writer_mae.add_scalar('MAE', mae, epoch)
print(phase + ' epoch {}:loss {:.4f} with data size {}'.format(
epoch, epoch_loss, dataset_sizes[phase]))
# checkpoint model if has best val loss yet
if phase == 'val' and epoch_loss < best_loss:
best_loss = epoch_loss
best_epoch = epoch
if not fine_tune:
checkpoint(model, best_loss, epoch, LR, filename='checkpoint_best')
elif fine_tune and not regression:
checkpoint(model, best_loss, epoch, LR, filename='classification_checkpoint_best')
else:
checkpoint(model, best_loss, epoch, LR, filename='regression_checkpoint_best')
# log training and validation loss over each epoch
with open("results/log_train", 'a') as logfile:
logwriter = csv.writer(logfile, delimiter=',')
if epoch == 1:
logwriter.writerow(["epoch", "train_loss", "val_loss"])
logwriter.writerow([epoch, last_train_loss, epoch_loss])
# Save model after each epoch
# checkpoint(model, best_loss, epoch, LR, filename='checkpoint')
total_done += batch_size
if total_done % (100 * batch_size) == 0:
print("completed " + str(total_done) + " so far in epoch")
# print elapsed time from the beginning after each epoch
print('Training complete in {:.0f}m {:.0f}s'.format(
(time.time() - since) // 60, (time.time() - since) % 60))
# total time
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
# load best model weights to return
if not fine_tune:
checkpoint_best = torch.load('results/checkpoint_best')
elif fine_tune and not regression:
checkpoint_best = torch.load('results/classification_checkpoint_best')
else:
checkpoint_best = torch.load('results/regression_checkpoint_best')
model = checkpoint_best['model']
return model, best_epoch
def train_cnn(PATH_TO_IMAGES, LR, WEIGHT_DECAY, fine_tune=False, regression=False, freeze=False, adam=False,
initial_model_path=None, initial_brixia_model_path=None, weighted_cross_entropy_batchwise=False,
modification=None, weighted_cross_entropy=False):
"""
Train torchvision model to NIH data given high level hyperparameters.
Args:
PATH_TO_IMAGES: path to NIH images
LR: learning rate
WEIGHT_DECAY: weight decay parameter for SGD
Returns:
preds: torchvision model predictions on test fold with ground truth for comparison
aucs: AUCs for each train,test tuple
"""
NUM_EPOCHS = 100
BATCH_SIZE = 32
try:
rmtree('results/')
except BaseException:
pass # directory doesn't yet exist, no need to clear it
os.makedirs("results/")
# use imagenet mean,std for normalization
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
N_LABELS = 14 # we are predicting 14 labels
N_COVID_LABELS = 3 # we are predicting 3 COVID labels
# define torchvision transforms
data_transforms = {
'train': transforms.Compose([
# transforms.RandomHorizontalFlip(),
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]),
'val': transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]),
}
# create train/val dataloaders
transformed_datasets = {}
transformed_datasets['train'] = CXR.CXRDataset(
path_to_images=PATH_TO_IMAGES,
fold='train',
transform=data_transforms['train'],
fine_tune=fine_tune,
regression=regression)
transformed_datasets['val'] = CXR.CXRDataset(
path_to_images=PATH_TO_IMAGES,
fold='val',
transform=data_transforms['val'],
fine_tune=fine_tune,
regression=regression)
dataloaders = {}
dataloaders['train'] = torch.utils.data.DataLoader(
transformed_datasets['train'],
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=8)
dataloaders['val'] = torch.utils.data.DataLoader(
transformed_datasets['val'],
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=8)
# please do not attempt to train without GPU as will take excessively long
if not use_gpu:
raise ValueError("Error, requires GPU")
if initial_model_path or initial_brixia_model_path:
if initial_model_path:
saved_model = torch.load(initial_model_path)
else:
saved_model = torch.load(initial_brixia_model_path)
model = saved_model['model']
del saved_model
if fine_tune and not initial_brixia_model_path:
num_ftrs = model.module.classifier.in_features
if freeze:
for feature in model.module.features:
for param in feature.parameters():
param.requires_grad = False
if feature == model.module.features.transition2:
break
if not regression:
model.module.classifier = nn.Linear(num_ftrs, N_COVID_LABELS)
else:
model.module.classifier = nn.Sequential(
nn.Linear(num_ftrs, 1),
nn.ReLU(inplace=True)
)
else:
model = models.densenet121(pretrained=True)
num_ftrs = model.classifier.in_features
model.classifier = nn.Linear(num_ftrs, N_LABELS)
if modification == 'transition_layer':
# num_ftrs = model.features.norm5.num_features
up1 = torch.nn.Sequential(torch.nn.ConvTranspose2d(num_ftrs, num_ftrs, kernel_size=3, stride=2, padding=1),
torch.nn.BatchNorm2d(num_ftrs),
torch.nn.ReLU(True))
up2 = torch.nn.Sequential(torch.nn.ConvTranspose2d(num_ftrs, num_ftrs, kernel_size=3, stride=2, padding=1),
torch.nn.BatchNorm2d(num_ftrs))
transition_layer = torch.nn.Sequential(up1, up2)
model.features.add_module('transition_chestX', transition_layer)
if modification == 'remove_last_block':
model.features.denseblock4 = nn.Sequential()
model.features.transition3 = nn.Sequential()
# model.features.norm5 = nn.BatchNorm2d(512)
# model.classifier = nn.Linear(512, N_LABELS)
if modification == 'remove_last_two_block':
model.features.denseblock4 = nn.Sequential()
model.features.transition3 = nn.Sequential()
model.features.transition2 = nn.Sequential()
model.features.denseblock3 = nn.Sequential()
model.features.norm5 = nn.BatchNorm2d(512)
model.classifier = nn.Linear(512, N_LABELS)
print(model)
# put model on GPU
if not initial_model_path:
model = nn.DataParallel(model)
model.to(device)
if regression:
criterion = nn.MSELoss()
else:
if weighted_cross_entropy:
pos_weights = transformed_datasets['train'].pos_neg_balance_weights()
print(pos_weights)
# pos_weights[pos_weights>40] = 40
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weights)
else:
criterion = nn.BCEWithLogitsLoss()
if adam:
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LR, weight_decay=WEIGHT_DECAY)
else:
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=LR, weight_decay=WEIGHT_DECAY, momentum=0.9)
dataset_sizes = {x: len(transformed_datasets[x]) for x in ['train', 'val']}
# train model
if regression:
model, best_epoch = train_model(model, criterion, optimizer, LR, num_epochs=NUM_EPOCHS,
dataloaders=dataloaders, dataset_sizes=dataset_sizes,
weight_decay=WEIGHT_DECAY, fine_tune=fine_tune, regression=regression)
else:
model, best_epoch = train_model(model, criterion, optimizer, LR, num_epochs=NUM_EPOCHS,
dataloaders=dataloaders, dataset_sizes=dataset_sizes, weight_decay=WEIGHT_DECAY,
weighted_cross_entropy_batchwise=weighted_cross_entropy_batchwise,
fine_tune=fine_tune)
# get preds and AUCs on test fold
preds, aucs = E.make_pred_multilabel(dataloaders['val'], model, save_as_csv=False, fine_tune=fine_tune)
return preds, aucs
from __future__ import print_function, division
# pytorch imports
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torchvision import transforms, utils
from tensorboardX import SummaryWriter
# general imports
import os
import time
from shutil import rmtree
# data science imports
import csv
import cxr_dataset as CXR
import eval_model as E
use_gpu = torch.cuda.is_available()
gpu_count = torch.cuda.device_count()
print("Available GPU count:" + str(gpu_count))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def checkpoint(model, best_loss, epoch, LR, filename):
"""
Saves checkpoint of torchvision model during training.
Args:
model: torchvision model to be saved
best_loss: best val loss achieved so far in training
epoch: current epoch of training
LR: current learning rate in training
Returns:
None
"""
print('saving')
state = {
'model': model,
'best_loss': best_loss,
'epoch': epoch,
'rng_state': torch.get_rng_state(),
'LR': LR
}
torch.save(state, 'results/' + filename)
def pos_neg_weights_in_batch(labels_batch):
num_total = labels_batch.shape[0] * labels_batch.shape[1]
num_positives = labels_batch.sum()
num_negatives = num_total - num_positives
if not num_positives == 0:
beta_p = num_negatives / num_positives
else:
beta_p = num_negatives
beta_p = torch.as_tensor(beta_p)
beta_p = beta_p.to(device)
beta_p = beta_p.type(torch.cuda.FloatTensor)
return beta_p
def train_model(
model,
criterion,
optimizer,
LR,
num_epochs,
dataloaders,
dataset_sizes,
weight_decay,
weighted_cross_entropy_batchwise=False,
fine_tune=False,
regression=False):
"""
Fine tunes torchvision model to NIH CXR data.
Args:
model: torchvision model to be finetuned (densenet-121 in this case)
criterion: loss criterion (binary cross entropy loss, BCELoss)
optimizer: optimizer to use in training (SGD)
LR: learning rate
num_epochs: continue training up to this many epochs
dataloaders: pytorch train and val dataloaders
dataset_sizes: length of train and val datasets
weight_decay: weight decay parameter we use in SGD with momentum
Returns:
model: trained torchvision model
best_epoch: epoch on which best model val loss was obtained
"""
since = time.time()
start_epoch = 1
best_loss = 999999
best_epoch = -1
last_train_loss = -1
tensorboard_writer_train = SummaryWriter('runs/loss/train_loss')
tensorboard_writer_val = SummaryWriter('runs/loss/val_loss')
if not fine_tune:
PRED_LABEL = [
'Atelectasis',
'Cardiomegaly',
'Effusion',
'Infiltration',
'Mass',
'Nodule',
'Pneumonia',
'Pneumothorax',
'Consolidation',
'Edema',
'Emphysema',
'Fibrosis',
'Pleural_Thickening',
'Hernia']
else:
PRED_LABEL = [
'Detector01',
'Detector2',
'Detector3']
if not regression:
tensorboard_writer_auc = {}
tensorboard_writer_AP = {}
for label in PRED_LABEL:
tensorboard_writer_auc[label] = SummaryWriter('runs/auc/'+label)
tensorboard_writer_AP[label] = SummaryWriter('runs/ap/' + label)
else:
tensorboard_writer_mae = SummaryWriter('runs/mae')
# iterate over epochs
for epoch in range(start_epoch, num_epochs + 1):
print('Epoch {}/{}'.format(epoch, num_epochs))
print('-' * 10)
# set model to train or eval mode based on whether we are in train or
# val; necessary to get correct predictions given batchnorm
for phase in ['train', 'val']:
if phase == 'train':
model.train(True)
else:
model.train(False)
running_loss = 0.0
total_done = 0
for data in dataloaders[phase]:
if not regression:
inputs, labels, _ = data
else:
inputs, ground_truths, _ = data
batch_size = inputs.shape[0]
inputs = inputs.to(device)
if not regression:
labels = (labels.to(device)).float()
else:
ground_truths = (ground_truths.to(device)).float()
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
# calculate gradient and update parameters in train phase
optimizer.zero_grad()
if weighted_cross_entropy_batchwise:
beta = pos_neg_weights_in_batch(labels)
criterion = nn.BCEWithLogitsLoss(pos_weight=beta)
if not regression:
loss = criterion(outputs, labels)
else:
ground_truths = ground_truths.unsqueeze(1)
loss = criterion(outputs, ground_truths)
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.item() * batch_size
epoch_loss = running_loss / dataset_sizes[phase]
if phase == 'train':
tensorboard_writer_train.add_scalar('Loss', epoch_loss, epoch)
last_train_loss = epoch_loss
elif phase == 'val':
tensorboard_writer_val.add_scalar('Loss', epoch_loss, epoch)
if not regression:
preds, aucs = E.make_pred_multilabel(dataloaders['val'], model, save_as_csv=False, fine_tune=fine_tune)
aucs.set_index('label', inplace=True)
print(aucs)
for label in PRED_LABEL:
tensorboard_writer_auc[label].add_scalar('AUC', aucs.loc[label, 'auc'], epoch)
tensorboard_writer_AP[label].add_scalar('AP', aucs.loc[label, 'AP'], epoch)
else:
mae, _, _ = E.evaluate_mae(dataloaders['val'], model)
print('MAE: ', mae)
tensorboard_writer_mae.add_scalar('MAE', mae, epoch)
print(phase + ' epoch {}:loss {:.4f} with data size {}'.format(
epoch, epoch_loss, dataset_sizes[phase]))
# checkpoint model if has best val loss yet
if phase == 'val' and epoch_loss < best_loss:
best_loss = epoch_loss
best_epoch = epoch
if not fine_tune:
checkpoint(model, best_loss, epoch, LR, filename='checkpoint_best_l1')
elif fine_tune and not regression:
checkpoint(model, best_loss, epoch, LR, filename='classification_checkpoint_best')
else:
checkpoint(model, best_loss, epoch, LR, filename='regression_checkpoint_best_l1')
# log training and validation loss over each epoch
with open("results/log_train", 'a') as logfile:
logwriter = csv.writer(logfile, delimiter=',')
if epoch == 1:
logwriter.writerow(["epoch", "train_loss", "val_loss"])
logwriter.writerow([epoch, last_train_loss, epoch_loss])
# Save model after each epoch
# checkpoint(model, best_loss, epoch, LR, filename='checkpoint')
total_done += batch_size
if total_done % (100 * batch_size) == 0:
print("completed " + str(total_done) + " so far in epoch")
# print elapsed time from the beginning after each epoch
print('Training complete in {:.0f}m {:.0f}s'.format(
(time.time() - since) // 60, (time.time() - since) % 60))
# total time
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
# load best model weights to return
if not fine_tune:
checkpoint_best = torch.load('results/checkpoint_best_l1')
elif fine_tune and not regression:
checkpoint_best = torch.load('results/classification_checkpoint_best')
else:
checkpoint_best = torch.load('results/regression_checkpoint_best_l1')
model = checkpoint_best['model']
return model, best_epoch
def train_cnn(PATH_TO_IMAGES, LR, WEIGHT_DECAY, fine_tune=False, regression=False, freeze=False, adam=False,
initial_model_path=None, initial_brixia_model_path=None, weighted_cross_entropy_batchwise=False,
modification=None, weighted_cross_entropy=False):
"""
Train torchvision model to NIH data given high level hyperparameters.
Args:
PATH_TO_IMAGES: path to NIH images
LR: learning rate
WEIGHT_DECAY: weight decay parameter for SGD
Returns:
preds: torchvision model predictions on test fold with ground truth for comparison
aucs: AUCs for each train,test tuple
"""
NUM_EPOCHS = 100
BATCH_SIZE = 32
try:
rmtree('results/')
except BaseException:
pass # directory doesn't yet exist, no need to clear it
os.makedirs("results/")
# use imagenet mean,std for normalization
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
N_LABELS = 14 # we are predicting 14 labels
N_COVID_LABELS = 3 # we are predicting 3 COVID labels
# define torchvision transforms
data_transforms = {
'train': transforms.Compose([
# transforms.RandomHorizontalFlip(),
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]),
'val': transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]),
}
# create train/val dataloaders
transformed_datasets = {}
transformed_datasets['train'] = CXR.CXRDataset(
path_to_images=PATH_TO_IMAGES,
fold='train',
transform=data_transforms['train'],
fine_tune=fine_tune,
regression=regression)
transformed_datasets['val'] = CXR.CXRDataset(
path_to_images=PATH_TO_IMAGES,
fold='val',
transform=data_transforms['val'],
fine_tune=fine_tune,
regression=regression)
dataloaders = {}
dataloaders['train'] = torch.utils.data.DataLoader(
transformed_datasets['train'],
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=8)
dataloaders['val'] = torch.utils.data.DataLoader(
transformed_datasets['val'],
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=8)
# please do not attempt to train without GPU as will take excessively long
if not use_gpu:
raise ValueError("Error, requires GPU")
if initial_model_path or initial_brixia_model_path:
if initial_model_path:
saved_model = torch.load(initial_model_path)
else:
saved_model = torch.load(initial_brixia_model_path)
model = saved_model['model']
del saved_model
if fine_tune and not initial_brixia_model_path:
num_ftrs = model.module.classifier.in_features
if freeze:
for feature in model.module.features:
for param in feature.parameters():
param.requires_grad = False
if feature == model.module.features.transition2:
break
if not regression:
model.module.classifier = nn.Linear(num_ftrs, N_COVID_LABELS)
else:
model.module.classifier = nn.Sequential(
nn.Linear(num_ftrs, 1),
nn.ReLU(inplace=True)
)
else:
model = models.densenet121(pretrained=True)
num_ftrs = model.classifier.in_features
model.classifier = nn.Linear(num_ftrs, N_LABELS)
if modification == 'transition_layer':
# num_ftrs = model.features.norm5.num_features
up1 = torch.nn.Sequential(torch.nn.ConvTranspose2d(num_ftrs, num_ftrs, kernel_size=3, stride=2, padding=1),
torch.nn.BatchNorm2d(num_ftrs),
torch.nn.ReLU(True))
up2 = torch.nn.Sequential(torch.nn.ConvTranspose2d(num_ftrs, num_ftrs, kernel_size=3, stride=2, padding=1),
torch.nn.BatchNorm2d(num_ftrs))
transition_layer = torch.nn.Sequential(up1, up2)
model.features.add_module('transition_chestX', transition_layer)
if modification == 'remove_last_block':
model.features.denseblock4 = nn.Sequential()
model.features.transition3 = nn.Sequential()
# model.features.norm5 = nn.BatchNorm2d(512)
# model.classifier = nn.Linear(512, N_LABELS)
if modification == 'remove_last_two_block':
model.features.denseblock4 = nn.Sequential()
model.features.transition3 = nn.Sequential()
model.features.transition2 = nn.Sequential()
model.features.denseblock3 = nn.Sequential()
model.features.norm5 = nn.BatchNorm2d(512)
model.classifier = nn.Linear(512, N_LABELS)
print(model)
# put model on GPU
if not initial_model_path:
model = nn.DataParallel(model)
model.to(device)
if regression:
criterion = nn.L1Loss()
else:
if weighted_cross_entropy:
pos_weights = transformed_datasets['train'].pos_neg_balance_weights()
print(pos_weights)
# pos_weights[pos_weights>40] = 40
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weights)
else:
criterion = nn.BCEWithLogitsLoss()
if adam:
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LR, weight_decay=WEIGHT_DECAY)
else:
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=LR, weight_decay=WEIGHT_DECAY, momentum=0.9)
dataset_sizes = {x: len(transformed_datasets[x]) for x in ['train', 'val']}
# train model
if regression:
model, best_epoch = train_model(model, criterion, optimizer, LR, num_epochs=NUM_EPOCHS,
dataloaders=dataloaders, dataset_sizes=dataset_sizes,
weight_decay=WEIGHT_DECAY, fine_tune=fine_tune, regression=regression)
else:
model, best_epoch = train_model(model, criterion, optimizer, LR, num_epochs=NUM_EPOCHS,
dataloaders=dataloaders, dataset_sizes=dataset_sizes, weight_decay=WEIGHT_DECAY,
weighted_cross_entropy_batchwise=weighted_cross_entropy_batchwise,
fine_tune=fine_tune)
# get preds and AUCs on test fold
preds, aucs = E.make_pred_multilabel(dataloaders['val'], model, save_as_csv=False, fine_tune=fine_tune)
return preds, aucs