def __init__(self,
full_file_path="",
transform=None,
gpu=True,
resize=320,
normalize=False):
self.resize = resize
self.gpu = gpu
self.root_PATH = root_PATH
self.saved_model_PATH = saved_model_PATH
self.root_PATH_dataset = root_PATH_dataset
self.file_name = full_file_path.split('/')[-1]
self.file_path = full_file_path
self.device = torch.device('cuda') if gpu else torch.device('cpu')
self.head_combo = []
self.kwargs = {"Relative_Position": {}, "Jigsaw": {}, "Rotation": {}}
self.model = models.densenet121().to(
self.device) if not full_file_path.__contains__(
"CC_GAN") else modified_densenet.densenet121().to(self.device)
self.out_D = {"Relative_Position": 8, "Rotation": 4, "Jigsaw": None}
self.transform_after_patching = transform
self.heads = None
self.normalize = normalize
def init_model(MODEL_NAME, N_LABELS):
# Check model used
import modified_densenet
import modified_alexnet
model = (
models.densenet121(pretrained=True) if MODEL_NAME == 'densenet' else
modified_densenet.densenet121(type=MODEL_NAME, pretrained=True)
if MODEL_NAME == 'va-densenet' or MODEL_NAME == 'reva-densenet'
or MODEL_NAME == 'fp-densenet' or MODEL_NAME == 'start-densenet'
or MODEL_NAME == 'every-densenet' or MODEL_NAME == 'sedensenet'
or MODEL_NAME == 'triplelossdensenet' else models.alexnet(
pretrained=True) if MODEL_NAME == 'alexnet' else modified_alexnet.
alexnet(type=MODEL_NAME, pretrained=True) if MODEL_NAME == 'va-alexnet'
or MODEL_NAME == 'reva-alexnet' or MODEL_NAME == 'fp-alexnet'
or MODEL_NAME == 'start-alexnet' else models.resnet152(
pretrained=True) if MODEL_NAME == 'resnet' else models.vgg16(
pretrained=True) if MODEL_NAME == 'VGG' else models.vgg16_bn(
pretrained=True) if MODEL_NAME == 'VGG_Bn' else '')
# get num_ftrs based on model name
num_ftrs = (model.classifier.in_features
if MODEL_NAME == 'densenet' or MODEL_NAME == 'va-densenet'
or MODEL_NAME == 'reva-densenet' or MODEL_NAME == 'fp-densenet'
or MODEL_NAME == 'start-densenet'
or MODEL_NAME == 'every-densenet' or MODEL_NAME == 'sedensenet'
or MODEL_NAME == 'triplelossdensenet' else
model.classifier[6].in_features if MODEL_NAME == 'alexnet'
or MODEL_NAME == 'va-alexnet' or MODEL_NAME == 'reva-alexnet'
or MODEL_NAME == 'fp-alexnet' or MODEL_NAME == 'start-alexnet'
or MODEL_NAME == 'VGG' or MODEL_NAME == 'VGG_Bn' else
model.fc.in_features if MODEL_NAME == 'resnet' else
model.fc3.in_features if MODEL_NAME == 'small_va' else '')
# change classifier class to N_LABELS
if (MODEL_NAME == 'densenet' or MODEL_NAME == 'va-densenet'
or MODEL_NAME == 'reva-densenet' or MODEL_NAME == 'fp-densenet'
or MODEL_NAME == 'start-densenet' or MODEL_NAME == 'every-densenet'
or MODEL_NAME == 'sedensenet'
or MODEL_NAME == 'triplelossdensenet'):
model.classifier = nn.Linear(num_ftrs, N_LABELS)
elif (MODEL_NAME == 'alexnet' or MODEL_NAME == 'va-alexnet'
or MODEL_NAME == 'va-alexnet' or MODEL_NAME == 'reva-alexnet'
or MODEL_NAME == 'fp-alexnet' or MODEL_NAME == 'start-alexnet'
or MODEL_NAME == 'VGG' or MODEL_NAME == 'VGG_Bn'):
model.classifier[6] = nn.Linear(num_ftrs, N_LABELS)
elif (MODEL_NAME == 'resnet'):
model.fc = nn.Linear(num_ftrs, N_LABELS)
else:
raise ValueError("Error model name")
return model
def __init__(self, pretrained=False):
"""
In the constructor we instantiate two nn.Linear modules and assign them as
member variables.
"""
super(Discriminator, self).__init__()
#Densenet with relu s replaced with leaky relu
self.discriminator = modified_densenet.densenet121(
num_classes=6, pretrained=pretrained)
if not pretrained:
for m in self.discriminator.modules():
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.Linear)):
nn.init.xavier_uniform_(
m.weight, gain=nn.init.calculate_gain('relu'))
def __init__(self):
super(mgn_dense, self).__init__()
model_ft = densenet121(pretrained=True)
block4 = copy.deepcopy(model_ft.features.denseblock4)
transition3 = copy.deepcopy(model_ft.features.transition3)
norm5 = copy.deepcopy(model_ft.features.norm5)
model_ft.features.transition3 = nn.Sequential()
model_ft.features.denseblock4 = nn.Sequential()
model_ft.features.norm5 = nn.Sequential()
self.base_mode = model_ft
############################## global
self.global_transition3 = copy.deepcopy(transition3)
self.global_block4 = copy.deepcopy(block4)
self.global_norm5 = copy.deepcopy(norm5)
self.global_avg = nn.AdaptiveAvgPool2d((1, 1))
self.global_reduce = nn.Conv2d(in_channels=1024,
out_channels=256,
kernel_size=1,
bias=False)
init.normal_(self.global_reduce.weight, std=0.01)
self.global_norm = nn.BatchNorm2d(num_features=256)
init.normal_(self.global_norm.weight, std=0.01)
self.global_prelu = nn.PReLU(256, 0.25)
self.global_fc = nn.Linear(in_features=1024, out_features=256)
init.normal_(self.global_fc.weight, std=0.01)
self.global_bottle = nn.Sequential(nn.BatchNorm1d(num_features=256),
nn.PReLU(256, 0.25),
nn.Dropout(0.5))
init.normal_(self.global_bottle[0].weight, std=0.01)
self.global_1024_linear = nn.Linear(in_features=256,
out_features=751,
bias=False)
init.normal_(self.global_1024_linear.weight, std=0.01)
############################## part2
self.part2_transition3 = copy.deepcopy(transition3)
self.part2_transition3.pool = nn.AdaptiveAvgPool2d((24, 8))
self.part2_block4 = copy.deepcopy(block4)
self.part2_norm5 = copy.deepcopy(norm5)
self.part2_avg = nn.AdaptiveAvgPool2d((1, 1))
self.part2_reduce = copy.deepcopy(self.global_reduce)
self.part2_norm = copy.deepcopy(self.global_norm)
self.part2_prelu = copy.deepcopy(self.global_prelu)
self.part2_fc = copy.deepcopy(self.global_fc)
self.part2_bottle = copy.deepcopy(self.global_bottle)
self.part2_1024 = copy.deepcopy(self.global_1024_linear)
self.part2_1_transition3 = copy.deepcopy(transition3)
self.part2_1_transition3.pool = nn.AdaptiveAvgPool2d((24, 8))
self.part2_1_block4 = copy.deepcopy(block4)
self.part2_1_norm5 = copy.deepcopy(norm5)
self.part2_1_avg = nn.AdaptiveAvgPool2d((1, 1))
self.part2_1_reduce = copy.deepcopy(self.global_reduce)
self.part2_1_bottle = copy.deepcopy(self.global_bottle)
self.part2_1_linear = copy.deepcopy(self.global_1024_linear)
self.part2_2_transition3 = copy.deepcopy(transition3)
self.part2_2_transition3.pool = nn.AdaptiveAvgPool2d((24, 8))
self.part2_2_block4 = copy.deepcopy(block4)
self.part2_2_norm5 = copy.deepcopy(norm5)
self.part2_2_avg = nn.AdaptiveAvgPool2d((1, 1))
self.part2_2_reduce = copy.deepcopy(self.global_reduce)
self.part2_2_bottle = copy.deepcopy(self.global_bottle)
self.part2_2_linear = copy.deepcopy(self.global_1024_linear)
############################## part3
self.part3_transition3 = copy.deepcopy(transition3)
self.part3_transition3.pool = nn.AdaptiveAvgPool2d((24, 8))
self.part3_block4 = copy.deepcopy(block4)
self.part3_norm5 = copy.deepcopy(norm5)
self.part3_avg = nn.AdaptiveAvgPool2d((1, 1))
self.part3_reduce = copy.deepcopy(self.global_reduce)
self.part3_norm = copy.deepcopy(self.global_norm)
self.part3_prelu = copy.deepcopy(self.global_prelu)
self.part3_fc = copy.deepcopy(self.global_fc)
self.part3_bottle = copy.deepcopy(self.global_bottle)
self.part3_1024 = copy.deepcopy(self.global_1024_linear)
self.part3_1_transition3 = copy.deepcopy(transition3)
self.part3_1_transition3.pool = nn.AdaptiveAvgPool2d((24, 8))
self.part3_1_block4 = copy.deepcopy(block4)
self.part3_1_norm5 = copy.deepcopy(norm5)
self.part3_1_avg = nn.AdaptiveAvgPool2d((1, 1))
self.part3_1_reduce = copy.deepcopy(self.global_reduce)
self.part3_1_bottle = copy.deepcopy(self.global_bottle)
self.part3_1_linear = copy.deepcopy(self.global_1024_linear)
self.part3_2_transition3 = copy.deepcopy(transition3)
self.part3_2_transition3.pool = nn.AdaptiveAvgPool2d((24, 8))
self.part3_2_block4 = copy.deepcopy(block4)
self.part3_2_norm5 = copy.deepcopy(norm5)
self.part3_2_avg = nn.AdaptiveAvgPool2d((1, 1))
self.part3_2_reduce = copy.deepcopy(self.global_reduce)
self.part3_2_bottle = copy.deepcopy(self.global_bottle)
self.part3_2_linear = copy.deepcopy(self.global_1024_linear)
self.part3_3_transition3 = copy.deepcopy(transition3)
self.part3_3_transition3.pool = nn.AdaptiveAvgPool2d((24, 8))
self.part3_3_block4 = copy.deepcopy(block4)
self.part3_3_norm5 = copy.deepcopy(norm5)
self.part3_3_avg = nn.AdaptiveAvgPool2d((1, 1))
self.part3_3_reduce = copy.deepcopy(self.global_reduce)
self.part3_3_bottle = copy.deepcopy(self.global_bottle)
self.part3_3_linear = copy.deepcopy(self.global_1024_linear)
def train_cnn(MODEL_NAME,
PRETRAINED,
FREEZE,
EPOCHS,
BATCH_SIZE,
N_LABELS,
OPTIMIZERS,
PATH_TO_IMAGES,
LR,
WEIGHT_DECAY,
LR_DECAY_STEPS,
DEBUG_MODE,
CHECKPOINT_PATH='',
DISTILLATE_WITH=''):
"""
Train torchvision model to NIH data given high level hyperparameters.
Args:
MODEL_NAME: model name
PRETRAINED: if model pretrained
FREEZE: model layer frozen or not
EPOCHS: epochs iteration
BATCH_SIZE: number of batch data per training
N_LABELS: number of class labels
OPTIMIZERS: optimizers used
PATH_TO_IMAGES: path to NIH images
LR: learning rate
WEIGHT_DECAY: weight decay parameter for SGD
LR_DECAY_STEPS: how many steps before LR decayed and dropped
DEBUG_MODE: if true then no log will be created
CHECKPOINT_PATH: load checkpoint path
DISTILLATE_WITH: distillate the model with
Returns:
# preds: torchvision model predictions on test fold with ground truth for comparison
# aucs: AUCs for each train,test tuple
"""
if not os.path.exists('results'):
os.makedirs('results')
# use imagenet mean,std for normalization
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = transforms.Normalize(mean=mean, std=std)
# define torchvision transforms
data_transforms = {
'train':
transforms.Compose([
transforms.Resize(256),
transforms.TenCrop(224),
transforms.Lambda(lambda crops: torch.stack(
[normalize(transforms.ToTensor()(crop)) for crop in crops]))
]),
'val':
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]),
'test':
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
}
# create train/val dataloaders
transformed_datasets = {
x: datasets.ImageFolder(os.path.join(PATH_TO_IMAGES, x),
data_transforms[x])
for x in ['train', 'val', 'test']
}
dataloaders = {
x: torch.utils.data.DataLoader(transformed_datasets[x],
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=0)
for x in ['train', 'val']
}
# please do not attempt to train without GPU as will take excessively long
if not use_gpu:
raise ValueError("Error, requires GPU")
# Check model used
import modified_densenet
import modified_alexnet
model = (
models.densenet121(
pretrained=PRETRAINED) if MODEL_NAME == 'densenet' else
modified_densenet.densenet121(type=MODEL_NAME, pretrained=PRETRAINED)
if MODEL_NAME == 'va-densenet' or MODEL_NAME == 'reva-densenet'
or MODEL_NAME == 'fp-densenet' or MODEL_NAME == 'start-densenet'
or MODEL_NAME == 'every-densenet' or MODEL_NAME == 'sedensenet'
or MODEL_NAME == 'triplelossdensenet' else models.alexnet(
pretrained=PRETRAINED) if MODEL_NAME == 'alexnet' else
modified_alexnet.alexnet(type=MODEL_NAME, pretrained=PRETRAINED)
if MODEL_NAME == 'va-alexnet' or MODEL_NAME == 'reva-alexnet'
or MODEL_NAME == 'fp-alexnet' or MODEL_NAME == 'start-alexnet' else
models.resnet152(pretrained=PRETRAINED) if MODEL_NAME ==
'resnet' else models.vgg16(
pretrained=PRETRAINED) if MODEL_NAME == 'VGG' else models.vgg16_bn(
pretrained=PRETRAINED) if MODEL_NAME == 'VGG_Bn' else '')
# get num_ftrs based on model name
num_ftrs = (model.classifier.in_features
if MODEL_NAME == 'densenet' or MODEL_NAME == 'va-densenet'
or MODEL_NAME == 'reva-densenet' or MODEL_NAME == 'fp-densenet'
or MODEL_NAME == 'start-densenet'
or MODEL_NAME == 'every-densenet' or MODEL_NAME == 'sedensenet'
or MODEL_NAME == 'triplelossdensenet' else
model.classifier[6].in_features if MODEL_NAME == 'alexnet'
or MODEL_NAME == 'va-alexnet' or MODEL_NAME == 'reva-alexnet'
or MODEL_NAME == 'fp-alexnet' or MODEL_NAME == 'start-alexnet'
or MODEL_NAME == 'VGG' or MODEL_NAME == 'VGG_Bn' else
model.fc.in_features if MODEL_NAME == 'resnet' else
model.fc3.in_features if MODEL_NAME == 'small_va' else '')
# change classifier class to N_LABELS
if (MODEL_NAME == 'densenet' or MODEL_NAME == 'va-densenet'
or MODEL_NAME == 'reva-densenet' or MODEL_NAME == 'fp-densenet'
or MODEL_NAME == 'start-densenet' or MODEL_NAME == 'every-densenet'
or MODEL_NAME == 'sedensenet'
or MODEL_NAME == 'triplelossdensenet'):
model.classifier = nn.Linear(num_ftrs, N_LABELS)
elif (MODEL_NAME == 'alexnet' or MODEL_NAME == 'va-alexnet'
or MODEL_NAME == 'va-alexnet' or MODEL_NAME == 'reva-alexnet'
or MODEL_NAME == 'fp-alexnet' or MODEL_NAME == 'start-alexnet'
or MODEL_NAME == 'VGG' or MODEL_NAME == 'VGG_Bn'):
model.classifier[6] = nn.Linear(num_ftrs, N_LABELS)
elif (MODEL_NAME == 'resnet'):
model.fc = nn.Linear(num_ftrs, N_LABELS)
else:
raise ValueError("Error model name")
if CHECKPOINT_PATH != '':
model = load_checkpoint(model, CHECKPOINT_PATH)
# show params to learn
if FREEZE:
for name, param in model.named_parameters():
attention_pattern = re.compile(
r'^(conv2d1x1|valinear|transconv|start|every|se_).+$')
classifier_pattern = re.compile(
r'^(classifier(?!\.\d)|classifier\.6|fc).+$')
if attention_pattern.match(name):
param.requires_grad = True
elif classifier_pattern.match(name) and CHECKPOINT_PATH == '':
param.requires_grad = True
else:
param.requires_grad = False
if FREEZE:
print('Params to learn:')
for name, param in model.named_parameters():
if param.requires_grad == True:
print(name)
print('==================================')
# Distillate
distillate_time = ''
if DISTILLATE_WITH != '':
print(f'Distillate with {DISTILLATE_WITH}')
distillate_time = datetime.datetime.now().isoformat()
model_distillate = models.densenet121(pretrained=PRETRAINED)
num_ftrs_distillate = model_distillate.classifier.in_features
model_distillate.classifier = nn.Linear(num_ftrs_distillate, N_LABELS)
model_distillate = load_checkpoint(model_distillate, DISTILLATE_WITH)
print('Loaded checkpoint for distillation')
model_distillate = model_distillate.cuda()
loading_bar = ''
dataloaders_length = len(dataloaders['train']) + len(
dataloaders['val'])
for i in range(dataloaders_length):
loading_bar += '-'
for phase in ['train', 'val']:
for data in dataloaders[phase]:
loading_bar = f'={loading_bar}'
loading_bar = loading_bar[:dataloaders_length]
print(f'Distillating: {loading_bar}', end='\r')
inputs, labels = data
if phase == 'train':
for i in range(10):
inp = inputs.clone()[:, i]
inp = inp.cuda()
labels = labels.cuda()
outputs = model_distillate(inp).cpu().data.numpy()
if len(outputs) != BATCH_SIZE:
outputs_padding = np.zeros((BATCH_SIZE, N_LABELS))
outputs_padding[:outputs.shape[0], :outputs.
shape[1]] = outputs
outputs = outputs_padding
if Path(f'results_distillation/d-{distillate_time}.npy'
).exists():
loaded_np = np.load(
f'results_distillation/d-{distillate_time}.npy'
)
outputs = np.append(loaded_np, [outputs], axis=0)
else:
outputs = [outputs]
np.save(
f'results_distillation/d-{distillate_time}.npy',
outputs)
else:
inputs = inputs.cuda()
labels = labels.cuda()
outputs = model_distillate(inputs).cpu().data.numpy()
if len(outputs) != BATCH_SIZE:
outputs_padding = np.zeros((BATCH_SIZE, N_LABELS))
outputs_padding[:outputs.shape[0], :outputs.
shape[1]] = outputs
outputs = outputs_padding
loaded_np = np.load(
f'results_distillation/d-{distillate_time}.npy')
outputs = np.append(loaded_np, [outputs], axis=0)
np.save(f'results_distillation/d-{distillate_time}.npy',
outputs)
print('')
# put model on GPU
model = model.cuda()
model.name = MODEL_NAME
# define criterion, optimizer for training
if distillate_time != '':
criterion = nn.MSELoss()
else:
# class_weight_value = [
# 1 - (1 / 77),
# 1 - (3 / 77),
# 1 - (12 / 77),
# 1 - (19 / 77),
# 1 - (3 / 77),
# 1 - (33 / 77),
# 1 - (6 / 77)
# ]
# class_weight = torch.FloatTensor(class_weight_value).cuda()
# criterion = nn.CrossEntropyLoss(weight=class_weight)
criterion = nn.CrossEntropyLoss()
# Check if SGD or Adam
optimizer = (optim.SGD(model.parameters(), lr=LR, momentum=0.9) if
OPTIMIZERS == 'SGD' else optim.Adam(model.parameters(), lr=LR)
if OPTIMIZERS == 'Adam' else '')
scheduler = lr_scheduler.StepLR(optimizer,
step_size=LR_DECAY_STEPS,
gamma=0.1)
dataset_sizes = {x: len(transformed_datasets[x]) for x in ['train', 'val']}
# train model
model, best_epoch = train_model(model,
criterion,
optimizer,
optim_name=OPTIMIZERS,
LR=LR,
num_epochs=EPOCHS,
dataloaders=dataloaders,
dataset_sizes=dataset_sizes,
weight_decay=WEIGHT_DECAY,
scheduler=scheduler,
debug_mode=DEBUG_MODE,
pretrained=PRETRAINED,
freeze=FREEZE,
checkpoint=CHECKPOINT_PATH,
distillate_time=distillate_time)
print("Finished Training")