def __init__(self, pathModel, nnArchitecture, nnClassCount, transCrop):
#---- Initialize the network
if nnArchitecture == 'DENSE-NET-121':
model = DenseNet121(nnClassCount, True).cuda()
elif nnArchitecture == 'DENSE-NET-169':
model = DenseNet169(nnClassCount, True).cuda()
elif nnArchitecture == 'DENSE-NET-201':
model = DenseNet201(nnClassCount, True).cuda()
model = torch.nn.DataParallel(model).cuda()
modelCheckpoint = torch.load(pathModel)
state_dict = modelCheckpoint['state_dict']
remove_data_parallel = False # Change if you don't want to use nn.DataParallel(model)
print('starging this')
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$'
)
for key in list(state_dict.keys()):
match = pattern.match(key)
new_key = match.group(1) + match.group(2) if match else key
new_key = new_key[7:] if remove_data_parallel else new_key
state_dict[new_key] = state_dict[key]
# Delete old key only if modified.
if match or remove_data_parallel:
del state_dict[key]
print('done this')
# if os.path.isfile(CKPT_PATH):
# print("=> loading checkpoint")
# checkpoint = torch.load(CKPT_PATH)
model.load_state_dict(modelCheckpoint['state_dict'])
print("=> loaded checkpoint")
self.model2 = model
self.model = model.module.densenet121.features
self.model.eval()
#---- Initialize the weights
self.weights = list(self.model.parameters())[-2]
#---- Initialize the image transform - resize + normalize
normalize = transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
transformList = []
transformList.append(transforms.Resize(transCrop))
transformList.append(transforms.ToTensor())
transformList.append(normalize)
self.transformSequence = transforms.Compose(transformList)
#added from chexnettrainer
#-------------------- SETTINGS: DATASET BUILDERS
transformList2 = []
transResize = 256
transformList2.append(transforms.Resize(transResize))
transformList2.append(transforms.TenCrop(224))
transformList2.append(
transforms.Lambda(lambda crops: torch.stack(
[transforms.ToTensor()(crop) for crop in crops])))
transformList2.append(
transforms.Lambda(lambda crops: torch.stack(
[normalize(crop) for crop in crops])))
self.transformSequence2 = transforms.Compose(transformList2)
def test (images, labels, pathModel, nnArchitecture, nnClassCount, nnIsTrained, trBatchSize, transResize, transCrop, launchTimeStamp):
CLASS_NAMES = [ 'A', 'B', 'C', 'D', 'E', '']
cudnn.benchmark = True
#-------------------- SETTINGS: NETWORK ARCHITECTURE, MODEL LOAD
if nnArchitecture == 'DENSE-NET-121': model = DenseNet121(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'DENSE-NET-169': model = DenseNet169(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'DENSE-NET-201': model = DenseNet201(nnClassCount, nnIsTrained).cuda()
model = torch.nn.DataParallel(model).cuda()
modelCheckpoint = torch.load(pathModel)
model.load_state_dict(modelCheckpoint['state_dict'])
#-------------------- SETTINGS: DATA TRANSFORMS, TEN CROPS
normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
#-------------------- SETTINGS: DATASET BUILDERS
transformList = []
transformList.append(transforms.Resize(transResize))
transformList.append(transforms.TenCrop(transCrop))
transformList.append(transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])))
transformList.append(transforms.Lambda(lambda crops: torch.stack([normalize(crop) for crop in crops])))
transformSequence=transforms.Compose(transformList)
datasetTest = ImageGenerator(images=images, labels=labels, transform=transformSequence)
dataLoaderTest = DataLoader(dataset=datasetTest, batch_size=trBatchSize, num_workers=8, shuffle=False, pin_memory=True)
outGT = torch.FloatTensor().cuda()
outPRED = torch.FloatTensor().cuda()
model.eval()
for i, (input, target) in enumerate(dataLoaderTest):
#target = target.cuda()
#outGT = torch.cat((outGT, target), 0)
bs, n_crops, c, h, w = input.size()
varInput = torch.autograd.Variable(input.view(-1, c, h, w).cuda(), volatile=True)
out = model(varInput)
outMean = out.view(bs, n_crops, -1).mean(1)
outPRED = torch.cat((outPRED, outMean.data), 0)
'''
aurocIndividual, cm = DensenetTrainer.computeAUROC(outGT, outPRED, nnClassCount, datasetTest)
aurocMean = np.array(aurocIndividual).mean()
print ('AUROC mean ', aurocMean)
for i in range (0, len(aurocIndividual)):
print (CLASS_NAMES[i], ' ', aurocIndividual[i])
print(cm)
'''
return outPRED
def train (pathDirData, pathFileTrain, pathFileVal, nnArchitecture, nnIsTrained, nnClassCount, trBatchSize, trMaxEpoch, transResize, transCrop, launchTimestamp, checkpoint):
#-------------------- SETTINGS: NETWORK ARCHITECTURE
if nnArchitecture == 'DENSE-NET-121': model = DenseNet121(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'DENSE-NET-169': model = DenseNet169(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'DENSE-NET-201': model = DenseNet201(nnClassCount, nnIsTrained).cuda()
model = torch.nn.DataParallel(model).cuda()
#-------------------- SETTINGS: DATA TRANSFORMS
normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
transformList = []
transformList.append(transforms.RandomResizedCrop(transCrop))
transformList.append(transforms.RandomHorizontalFlip())
transformList.append(transforms.RandomVerticalFlip())
transformList.append(transforms.RandomRotation(25))
transformList.append(transforms.ToTensor())
transformList.append(normalize)
transformSequence=transforms.Compose(transformList)
#-------------------- SETTINGS: DATASET BUILDERS
datasetTrain = DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileTrain, transform=transformSequence)
datasetVal = DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileVal, transform=transformSequence)
dataLoaderTrain = DataLoader(dataset=datasetTrain, batch_size=trBatchSize, shuffle=True, num_workers=24, pin_memory=True)
dataLoaderVal = DataLoader(dataset=datasetVal, batch_size=trBatchSize, shuffle=False, num_workers=24, pin_memory=True)
#-------------------- SETTINGS: OPTIMIZER & SCHEDULER
optimizer = optim.Adam (model.parameters(), lr=0.000001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-6)
scheduler = ReduceLROnPlateau(optimizer, factor = 0.1, patience = 3, mode = 'min')
#-------------------- SETTINGS: LOSS
#loss = torch.nn.BCELoss(size_average = True)
loss = torch.nn.BCELoss(size_average=True)
#---- Load checkpoint
if checkpoint != None:
modelCheckpoint = torch.load(checkpoint)
model.load_state_dict(modelCheckpoint['state_dict'])
optimizer.load_state_dict(modelCheckpoint['optimizer'])
#---- TRAIN THE NETWORK
lossMIN = 100000
for epochID in range (0, trMaxEpoch):
timestampTime = time.strftime("%H%M%S")
timestampDate = time.strftime("%d%m%Y")
timestampSTART = timestampDate + '-' + timestampTime
DensenetTrainer.epochTrain (model, dataLoaderTrain, optimizer, scheduler, trMaxEpoch, nnClassCount, loss)
lossVal, losstensor = DensenetTrainer.epochVal (model, dataLoaderVal, optimizer, scheduler, trMaxEpoch, nnClassCount, loss)
timestampTime = time.strftime("%H%M%S")
timestampDate = time.strftime("%d%m%Y")
timestampEND = timestampDate + '-' + timestampTime
scheduler.step(losstensor.data[0])
if lossVal < lossMIN:
lossMIN = lossVal
torch.save({'epoch': epochID + 1, 'state_dict': model.state_dict(), 'best_loss': lossMIN, 'optimizer' : optimizer.state_dict()}, 'm-' + launchTimestamp + '.pth.tar')
print ('Epoch [' + str(epochID + 1) + '] [save] [' + timestampEND + '] loss= ' + str(lossVal))
pathModel = 'm-' + launchTimestamp + '.pth.tar'
DensenetTrainer.test(pathDirData, pathFileVal, pathModel, nnArchitecture, nnClassCount, nnIsTrained, 16,
transResize, transCrop, launchTimestamp)
else:
print ('Epoch [' + str(epochID + 1) + '] [----] [' + timestampEND + '] loss= ' + str(lossVal))
def test (pathDirData, pathFileTest, pathModel, nnArchitecture, nnClassCount, nnIsTrained, trBatchSize, transResize, transCrop, launchTimeStamp):
CLASS_NAMES = [ 'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass', 'Nodule', 'Pneumonia',
'Pneumothorax', 'Consolidation', 'Edema', 'Emphysema', 'Fibrosis', 'Pleural_Thickening', 'Hernia']
cudnn.benchmark = True
#-------------------- SETTINGS: NETWORK ARCHITECTURE, MODEL LOAD
if nnArchitecture == 'DENSE-NET-121': model = DenseNet121(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'DENSE-NET-169': model = DenseNet169(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'DENSE-NET-201': model = DenseNet201(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'RES-NET-50': model = ResNet50(nnClassCount, nnIsTrained).cuda()
#elif nnArchitecture == 'SE-RES-NET-50': model = SE_ResNet50(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'SE-RES-NET-50': model = senet.se_resnet50(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'SE-NET-154': model = senet.senet154(nnClassCount, nnIsTrained).cuda()
elif nnArchitecture == 'SE-DENSE-NET-121': model = SE_DenseNet121(nnClassCount, nnIsTrained).cuda()
model = torch.nn.DataParallel(model).cuda()
modelCheckpoint = torch.load(pathModel)
model.load_state_dict(modelCheckpoint['state_dict'])
#-------------------- SETTINGS: DATA TRANSFORMS, TEN CROPS
normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
#-------------------- SETTINGS: DATASET BUILDERS
transformList = []
transformList.append(transforms.Resize(transResize))
transformList.append(transforms.TenCrop(transCrop))
transformList.append(transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])))
transformList.append(transforms.Lambda(lambda crops: torch.stack([normalize(crop) for crop in crops])))
transformSequence=transforms.Compose(transformList)
datasetTest = DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileTest, transform=transformSequence)
dataLoaderTest = DataLoader(dataset=datasetTest, batch_size=trBatchSize, num_workers=8, shuffle=False, pin_memory=True)
outGT = torch.FloatTensor().cuda()
outPRED = torch.FloatTensor().cuda()
model.eval()
for i, (input, target) in enumerate(dataLoaderTest):
target = target.cuda()
outGT = torch.cat((outGT, target), 0)
bs, n_crops, c, h, w = input.size()
varInput = torch.autograd.Variable(input.view(-1, c, h, w).cuda(), volatile=True)
out = model(varInput)
outMean = out.view(bs, n_crops, -1).mean(1)
outPRED = torch.cat((outPRED, outMean.data), 0)
aurocIndividual = ChexnetTrainer.computeAUROC(outGT, outPRED, nnClassCount)
aurocMean = np.array(aurocIndividual).mean()
print ('AUROC mean ', aurocMean)
for i in range (0, len(aurocIndividual)):
print (CLASS_NAMES[i], ' ', aurocIndividual[i])
return
