# Base / Native
import math
import os
import pickle
import re
import warnings
warnings.filterwarnings('ignore')
# Numerical / Array
import lifelines
from lifelines.utils import concordance_index
from lifelines import CoxPHFitter
from lifelines.datasets import load_regression_dataset
from lifelines.utils import k_fold_cross_validation
from lifelines.statistics import logrank_test
from imblearn.over_sampling import RandomOverSampler
import numpy as np
# Torch
import torch
import torch.nn as nn
from torch.nn import init, Parameter
from torch.utils.data._utils.collate import *
from torch.utils.data.dataloader import default_collate
import torch_geometric
from torch_geometric.data import Batch
################
# Regularization
################
def regularize_weights(model, reg_type=None):
l1_reg = None
for W in model.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
return l1_reg
def regularize_path_weights(model, reg_type=None):
l1_reg = None
for W in model.classifier.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
for W in model.linear.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
return l1_reg
def regularize_MM_weights(model, reg_type=None):
l1_reg = None
if model.module.__hasattr__('omic_net'):
for W in model.module.omic_net.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_h_path'):
for W in model.module.linear_h_path.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_h_omic'):
for W in model.module.linear_h_omic.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_h_grph'):
for W in model.module.linear_h_grph.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_z_path'):
for W in model.module.linear_z_path.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_z_omic'):
for W in model.module.linear_z_omic.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_z_grph'):
for W in model.module.linear_z_grph.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_o_path'):
for W in model.module.linear_o_path.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_o_omic'):
for W in model.module.linear_o_omic.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('linear_o_grph'):
for W in model.module.linear_o_grph.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('encoder1'):
for W in model.module.encoder1.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('encoder2'):
for W in model.module.encoder2.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
if model.module.__hasattr__('classifier'):
for W in model.module.classifier.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
return l1_reg
def regularize_MM_omic(model, reg_type=None):
l1_reg = None
if model.module.__hasattr__('omic_net'):
for W in model.module.omic_net.parameters():
if l1_reg is None:
l1_reg = torch.abs(W).sum()
else:
l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)
return l1_reg
################
# Network Initialization
################
def init_weights(net, init_type='orthogonal', init_gain=0.02):
"""Initialize network weights.
Parameters:
net (network) -- network to be initialized
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
work better for some applications. Feel free to try yourself.
"""
def init_func(m): # define the initialization function
classname = m.__class__.__name__
if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
if init_type == 'normal':
init.normal_(m.weight.data, 0.0, init_gain)
elif init_type == 'xavier':
init.xavier_normal_(m.weight.data, gain=init_gain)
elif init_type == 'kaiming':
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
elif init_type == 'orthogonal':
init.orthogonal_(m.weight.data, gain=init_gain)
else:
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
if hasattr(m, 'bias') and m.bias is not None:
init.constant_(m.bias.data, 0.0)
elif classname.find('BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
init.normal_(m.weight.data, 1.0, init_gain)
init.constant_(m.bias.data, 0.0)
print('initialize network with %s' % init_type)
net.apply(init_func) # apply the initialization function <init_func>
def init_max_weights(module):
for m in module.modules():
if type(m) == nn.Linear:
stdv = 1. / math.sqrt(m.weight.size(1))
m.weight.data.normal_(0, stdv)
m.bias.data.zero_()
def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
"""Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
Parameters:
net (network) -- the network to be initialized
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
gain (float) -- scaling factor for normal, xavier and orthogonal.
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
Return an initialized network.
"""
if len(gpu_ids) > 0:
assert(torch.cuda.is_available())
net.to(gpu_ids[0])
net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs
if init_type != 'max' and init_type != 'none':
print("Init Type:", init_type)
init_weights(net, init_type, init_gain=init_gain)
elif init_type == 'none':
print("Init Type: Not initializing networks.")
elif init_type == 'max':
print("Init Type: Self-Normalizing Weights")
return net
################
# Freeze / Unfreeze
################
def unfreeze_unimodal(opt, model, epoch):
if opt.mode == 'graphomic':
if epoch == 5:
dfs_unfreeze(model.module.omic_net)
print("Unfreezing Omic")
if epoch == 5:
dfs_unfreeze(model.module.grph_net)
print("Unfreezing Graph")
elif opt.mode == 'pathomic':
if epoch == 5:
dfs_unfreeze(model.module.omic_net)
print("Unfreezing Omic")
elif opt.mode == 'pathgraph':
if epoch == 5:
dfs_unfreeze(model.module.grph_net)
print("Unfreezing Graph")
elif opt.mode == "pathgraphomic":
if epoch == 5:
dfs_unfreeze(model.module.omic_net)
print("Unfreezing Omic")
if epoch == 5:
dfs_unfreeze(model.module.grph_net)
print("Unfreezing Graph")
elif opt.mode == "omicomic":
if epoch == 5:
dfs_unfreeze(model.module.omic_net)
print("Unfreezing Omic")
elif opt.mode == "graphgraph":
if epoch == 5:
dfs_unfreeze(model.module.grph_net)
print("Unfreezing Graph")
def dfs_freeze(model):
for name, child in model.named_children():
for param in child.parameters():
param.requires_grad = False
dfs_freeze(child)
def dfs_unfreeze(model):
for name, child in model.named_children():
for param in child.parameters():
param.requires_grad = True
dfs_unfreeze(child)
def print_if_frozen(module):
for idx, child in enumerate(module.children()):
for param in child.parameters():
if param.requires_grad == True:
print("Learnable!!! %d:" % idx, child)
else:
print("Still Frozen %d:" % idx, child)
def unfreeze_vgg_features(model, epoch):
epoch_schedule = {30:45}
unfreeze_index = epoch_schedule[epoch]
for idx, child in enumerate(model.features.children()):
if idx > unfreeze_index:
print("Unfreezing %d:" %idx, child)
for param in child.parameters():
param.requires_grad = True
else:
print("Still Frozen %d:" %idx, child)
continue
################
# Collate Utils
################
def mixed_collate(batch):
elem = batch[0]
elem_type = type(elem)
transposed = zip(*batch)
return [Batch.from_data_list(samples, []) if type(samples[0]) is torch_geometric.data.data.Data else default_collate(samples) for samples in transposed]
################
# Survival Utils
################
def CoxLoss(survtime, censor, hazard_pred, device):
# This calculation credit to Travers Ching https://github.com/traversc/cox-nnet
# Cox-nnet: An artificial neural network method for prognosis prediction of high-throughput omics data
current_batch_len = len(survtime)
R_mat = np.zeros([current_batch_len, current_batch_len], dtype=int)
for i in range(current_batch_len):
for j in range(current_batch_len):
R_mat[i,j] = survtime[j] >= survtime[i]
R_mat = torch.FloatTensor(R_mat).to(device)
theta = hazard_pred.reshape(-1)
exp_theta = torch.exp(theta)
loss_cox = -torch.mean((theta - torch.log(torch.sum(exp_theta*R_mat, dim=1))) * censor)
return loss_cox
def accuracy(output, labels):
preds = output.max(1)[1].type_as(labels)
correct = preds.eq(labels).double()
correct = correct.sum()
return correct / len(labels)
def accuracy_cox(hazardsdata, labels):
# This accuracy is based on estimated survival events against true survival events
median = np.median(hazardsdata)
hazards_dichotomize = np.zeros([len(hazardsdata)], dtype=int)
hazards_dichotomize[hazardsdata > median] = 1
correct = np.sum(hazards_dichotomize == labels)
return correct / len(labels)
def cox_log_rank(hazardsdata, labels, survtime_all):
median = np.median(hazardsdata)
hazards_dichotomize = np.zeros([len(hazardsdata)], dtype=int)
hazards_dichotomize[hazardsdata > median] = 1
idx = hazards_dichotomize == 0
T1 = survtime_all[idx]
T2 = survtime_all[~idx]
E1 = labels[idx]
E2 = labels[~idx]
results = logrank_test(T1, T2, event_observed_A=E1, event_observed_B=E2)
pvalue_pred = results.p_value
return(pvalue_pred)
def CIndex(hazards, labels, survtime_all):
concord = 0.
total = 0.
N_test = labels.shape[0]
for i in range(N_test):
if labels[i] == 1:
for j in range(N_test):
if survtime_all[j] > survtime_all[i]:
total += 1
if hazards[j] < hazards[i]: concord += 1
elif hazards[j] < hazards[i]: concord += 0.5
return(concord/total)
def CIndex_lifeline(hazards, labels, survtime_all):
return(concordance_index(survtime_all, -hazards, labels))