├── scripts
    ├── extend_kd_tree_offline.py
    ├── Augmentation.py
    └── Training_new_augmentation.py
└── README.md


/scripts/extend_kd_tree_offline.py:
--------------------------------------------------------------------------------
  1 | import pandas as pd
  2 | import numpy as np
  3 | import os
  4 | from PIL import Image
  5 | import albumentations as A
  6 | from scipy.spatial import KDTree, cKDTree
  7 | import pickle
  8 | import matplotlib.pyplot as plt
  9 | import time
 10 | from tqdm import tqdm
 11 | import argparse
 12 | import sys, getopt
 13 | 
 14 | argv = sys.argv[1:]
 15 | 
 16 | 
 17 | #parser parameters
 18 | parser = argparse.ArgumentParser(description='Configurations to train models.')
 19 | parser.add_argument('-i', '--INPUT', help='input to extend',type=str, default='input.pickle')
 20 | parser.add_argument('-o', '--OUTPUT', help='where to store file',type=str, default='output.pickle')
 21 | parser.add_argument('-d', '--DATA_TO_ADD', help='csv_file',type=int, default=32)
 22 | 
 23 | 
 24 | args = parser.parse_args()
 25 | 
 26 | INPUT_DATA = args.DATA_TO_ADD
 27 | OUTPUT_FILE = args.OUTPUT
 28 | INPUT_FILE = args.INPUT
 29 | 
 30 | def H_E_Staining(img, Io=240, alpha=1, beta=0.15):
 31 | 
 32 | 	# define height and width of image
 33 | 	h, w, c = img.shape
 34 | 
 35 | 	# reshape image
 36 | 	img = img.reshape((-1,3))
 37 | 
 38 | 	# calculate optical density
 39 | 	OD = -np.log((img.astype(np.float)+1)/Io)
 40 | 
 41 | 	# remove transparent pixels
 42 | 	ODhat = OD[~np.any(OD<beta, axis=1)]
 43 | 
 44 | 	# compute eigenvectors
 45 | 	eigvals, eigvecs = np.linalg.eigh(np.cov(ODhat.T))
 46 | 
 47 | 	#eigvecs *= -1
 48 | 
 49 | 	#project on the plane spanned by the eigenvectors corresponding to the two 
 50 | 	# largest eigenvalues    
 51 | 	That = ODhat.dot(eigvecs[:,1:3])
 52 | 
 53 | 	phi = np.arctan2(That[:,1],That[:,0])
 54 | 
 55 | 	minPhi = np.percentile(phi, alpha)
 56 | 	maxPhi = np.percentile(phi, 100-alpha)
 57 | 
 58 | 	vMin = eigvecs[:,1:3].dot(np.array([(np.cos(minPhi), np.sin(minPhi))]).T)
 59 | 	vMax = eigvecs[:,1:3].dot(np.array([(np.cos(maxPhi), np.sin(maxPhi))]).T)
 60 | 
 61 | 	# a heuristic to make the vector corresponding to hematoxylin first and the 
 62 | 	# one corresponding to eosin second
 63 | 	if vMin[0] > vMax[0]:
 64 | 		HE = np.array((vMin[:,0], vMax[:,0])).T
 65 | 	else:
 66 | 		HE = np.array((vMax[:,0], vMin[:,0])).T
 67 | 
 68 | 	return HE
 69 | 
 70 | fname = INPUT_FILE
 71 | 
 72 | with open(fname, 'rb') as f:
 73 | 	kdtree = pickle.load(f)
 74 | 
 75 | HEs_general = kdtree.data
 76 | 
 77 | input_csv = INPUT_DATA
 78 | input_data = pd.read_csv(input_csv, sep=',',header=None).values
 79 | 
 80 | def extend_stains(kdtree, new_data, save_new_array = False, PERC=1.0):
 81 | 	
 82 | 	HEs_new_stains = []
 83 | 	
 84 | 	threshold_value = int(len(new_data)*PERC)
 85 | 	
 86 | 	i = 0
 87 | 		
 88 | 	HEs_general = kdtree.data
 89 | 	
 90 | 	np.random.shuffle(new_data)
 91 | 	
 92 | 	for i in tqdm(range(threshold_value)):
 93 | 
 94 | 		patch = new_data[i,0]
 95 | 		
 96 | 		img = Image.open(patch)
 97 | 		img_np = np.asarray(img)
 98 | 		
 99 | 		HE = H_E_Staining(img_np)
100 | 		
101 | 		HE = np.reshape(HE, 6)
102 | 		HEs_new_stains.append(HE)
103 | 		
104 | 		img.close()
105 | 		
106 | 		#i = i + 1
107 | 
108 | 	HEs_new_stains = np.array(HEs_new_stains)
109 | 	HEs_general = np.append(HEs_general,HEs_new_stains,axis=0)
110 | 	
111 | 	new_kdtree = cKDTree(HEs_general)
112 | 	
113 | 	if (save_new_array==True):
114 | 		
115 | 		fname = OUTPUT_FILE
116 | 		with open(fname, 'wb') as f:
117 | 			pickle.dump(kdtree, f)
118 | 		
119 | 	return new_kdtree
120 | 
121 | start_time = time.time()
122 | new_kdtree = extend_stains(kdtree, input_data, save_new_array = True, PERC=1.0)
123 | elapsed_time = time.time() - start_time
124 | print("elapsed time " + str(elapsed_time))


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Data_Driven_Color_Augmentation
 2 | 
 3 | Implementation of "Data-driven color augmentation for H&E stained images in computational pathology".
 4 | 
 5 | ## Reference
 6 | If you find this repository useful in your research, please cite:
 7 | 
 8 | [1] Marini N., Otálora S., Wodzinski M., Tomassini S. Dragoni A.F., Marchand-Maillet S., Dominguez P., Duran-Lopez L., Vatrano S., Müller H. & Atzori M., Data-driven color augmentation for H&E stained images in computational pathology.
 9 | 
10 | Paper link: https://www.sciencedirect.com/science/article/pii/S2153353922007830
11 | 
12 | ## Requirements
13 | Python==3.6.9, albumentations==0.1.8, numpy==1.17.3, opencv==4.2.0, pandas==0.25.2, pillow==6.1.0, torchvision==0.8.1, pytorch==1.7.0
14 | 
15 | ## CSV Input Files:
16 | CSV files are used as input for the scripts. For each partition (train, validation, test), the csv file has path_to_image, class_label as columns.
17 | For prostate experiments, the class_label can be: 
18 | 0: benign
19 | 1: Gleason pattern 3
20 | 2: Gleason pattern 4
21 | 3: Gleason pattern 5
22 | 
23 | For colon experiments, the class_label can be:
24 | 0: cancer
25 | 1: dysplasia
26 | 2: normal glands
27 | 
28 | ## Augmentation
29 | Methods to perform data drive color augmentation (Augmentation.py):
30 | - new_color_augmentation (HSC color augmentation):
31 |   * patch_np: numpy array for the input patch (224x224)
32 |   * kdtree: the database where acceptable color variations are stored
33 |   * alpha: neighbors
34 |   * beta: radius
35 |   * shift_value: perturbation to apply to Hue, Saturation, Contrast 
36 | - new_stain_augmentation (perturbation of H&E channels):
37 |   * patch_np: numpy array for the input patch (224x224)
38 |   * kdtree: the database where acceptable color variations are stored
39 |   * alpha: neighbors
40 |   * beta: radius
41 |   * sigma1: range (-sigma1, sigma1) to generate random value to multiply to H&E components
42 |   * sigma2: range (-sigma2, sigma2) to generate random value to add to H&E components
43 | 
44 | ## Database
45 | Database including color variations: https://zenodo.org/record/7505727#.Y7ayO3bMJPY.
46 | 
47 | Method to extend database with new histopathology patches (extend_kd_tree_offline)
48 |   * -i: input pickle file (database to extend)
49 |   * -o: output pickle file
50 |   * -d: csv including patches to extend database
51 | 
52 | ## Training
53 | Scripts to train the CNN at path-level, in a fully-supervised fashion.
54 | Some parameters must be manually changed, such as the number of classes (output of the network).
55 | 
56 | - Training_new_augmentation.py -n -b -c -e -f -i -o -a -d. The script is used to train the CNN without any augmentation (no_augment), with colour augmentation (augment).
57 |   * -n: number of the experiment for the training
58 |   * -b: batch size (32)
59 |   * -c: CNN backbone to use (densenet121)
60 |   * -e: number of epochs (10)
61 |   * -t: task of the network (no_augment, augment, normalizer)
62 |   * -f: if True an embedding layer with 128 nodes is inserted before the output layer
63 |   * -i: path of the folder where the input csvs for training (train.csv), validation (valid.csv) and testing (test.csv) are stored
64 |   * -o: path of the folder where to store the CNN’s weights.
65 |   * -a: new augmentation to use: color (HSC color augmentation), stain (H&E stain augmentation), he (H&E-adversarial CNN + HSC color augmentation)
66 |   * -x: extend (False): add color variations training data to color variation dataset
67 |   * -d: database: database including color variations
68 | 
69 | 
70 | ## Acknoledgements
71 | This project has received funding from the EuropeanUnion’s Horizon 2020 research and innovation programme under grant agree-ment No. 825292 [ExaMode](http://www.examode.eu). Infrastructure fromthe SURFsara HPC center was used to train the CNN models in parallel. 
72 | 


--------------------------------------------------------------------------------
/scripts/Augmentation.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | import pandas as pd
  3 | from PIL import Image
  4 | import albumentations as A
  5 | import warnings 
  6 | 
  7 | from scipy.spatial import KDTree, cKDTree
  8 | 
  9 | warnings.filterwarnings("ignore")
 10 | 
 11 | def H_E_Staining(img, Io=240, alpha=1, beta=0.15):
 12 | 
 13 | 	# define height and width of image
 14 | 	h, w, c = img.shape
 15 | 
 16 | 	# reshape image
 17 | 	img = img.reshape((-1,3))
 18 | 
 19 | 	# calculate optical density
 20 | 	OD = -np.log((img.astype(np.float)+1)/Io)
 21 | 
 22 | 	# remove transparent pixels
 23 | 	ODhat = OD[~np.any(OD<beta, axis=1)]
 24 | 
 25 | 	# compute eigenvectors
 26 | 	eigvals, eigvecs = np.linalg.eigh(np.cov(ODhat.T))
 27 | 
 28 | 	#eigvecs *= -1
 29 | 
 30 | 	#project on the plane spanned by the eigenvectors corresponding to the two 
 31 | 	# largest eigenvalues    
 32 | 	That = ODhat.dot(eigvecs[:,1:3])
 33 | 
 34 | 	phi = np.arctan2(That[:,1],That[:,0])
 35 | 
 36 | 	minPhi = np.percentile(phi, alpha)
 37 | 	maxPhi = np.percentile(phi, 100-alpha)
 38 | 
 39 | 	vMin = eigvecs[:,1:3].dot(np.array([(np.cos(minPhi), np.sin(minPhi))]).T)
 40 | 	vMax = eigvecs[:,1:3].dot(np.array([(np.cos(maxPhi), np.sin(maxPhi))]).T)
 41 | 
 42 | 	# a heuristic to make the vector corresponding to hematoxylin first and the 
 43 | 	# one corresponding to eosin second
 44 | 	if vMin[0] > vMax[0]:
 45 | 		HE = np.array((vMin[:,0], vMax[:,0])).T
 46 | 	else:
 47 | 		HE = np.array((vMax[:,0], vMin[:,0])).T
 48 | 
 49 | 	return HE
 50 | 
 51 | def unique_elements(array):
 52 | 	
 53 | 	_, counts = np.unique(array, return_counts=True)
 54 | 	
 55 | 	b = True
 56 | 	
 57 | 	for c in counts:
 58 | 		
 59 | 		if (c>1):
 60 | 			
 61 | 			b = False
 62 | 	
 63 | 	return b
 64 | 
 65 | def normalizeStaining(img, HERef, Io=240, alpha=1, beta=0.15):
 66 | 
 67 | 	maxCRef = np.array([1.9705, 1.0308])
 68 | 	
 69 | 	# define height and width of image
 70 | 	h, w, c = img.shape
 71 | 	
 72 | 	# reshape image
 73 | 	img = img.reshape((-1,3))
 74 | 
 75 | 	# calculate optical density
 76 | 	OD = -np.log((img.astype(np.float)+1)/Io)
 77 | 	
 78 | 	# remove transparent pixels
 79 | 	ODhat = OD[~np.any(OD<beta, axis=1)]
 80 | 		
 81 | 	# compute eigenvectors
 82 | 	eigvals, eigvecs = np.linalg.eigh(np.cov(ODhat.T))
 83 | 	
 84 | 	#eigvecs *= -1
 85 | 	
 86 | 	#project on the plane spanned by the eigenvectors corresponding to the two 
 87 | 	# largest eigenvalues    
 88 | 	That = ODhat.dot(eigvecs[:,1:3])
 89 | 	
 90 | 	phi = np.arctan2(That[:,1],That[:,0])
 91 | 	
 92 | 	minPhi = np.percentile(phi, alpha)
 93 | 	maxPhi = np.percentile(phi, 100-alpha)
 94 | 	
 95 | 	vMin = eigvecs[:,1:3].dot(np.array([(np.cos(minPhi), np.sin(minPhi))]).T)
 96 | 	vMax = eigvecs[:,1:3].dot(np.array([(np.cos(maxPhi), np.sin(maxPhi))]).T)
 97 | 	
 98 | 	# a heuristic to make the vector corresponding to hematoxylin first and the 
 99 | 	# one corresponding to eosin second
100 | 	if vMin[0] > vMax[0]:
101 | 		HE = np.array((vMin[:,0], vMax[:,0])).T
102 | 	else:
103 | 		HE = np.array((vMax[:,0], vMin[:,0])).T
104 | 	
105 | 	# rows correspond to channels (RGB), columns to OD values
106 | 	Y = np.reshape(OD, (-1, 3)).T
107 | 	
108 | 	# determine concentrations of the individual stains
109 | 	C = np.linalg.lstsq(HE,Y, rcond=None)[0]
110 | 	
111 | 	# normalize stain concentrations
112 | 	maxC = np.array([np.percentile(C[0,:], 99), np.percentile(C[1,:],99)])
113 | 	tmp = np.divide(maxC,maxCRef)
114 | 	C2 = np.divide(C,tmp[:, np.newaxis])
115 | 	
116 | 	# recreate the image using reference mixing matrix
117 | 	Inorm = np.multiply(Io, np.exp(-HERef.dot(C2)))
118 | 	Inorm[Inorm>255] = 254
119 | 	Inorm = np.reshape(Inorm.T, (h, w, 3)).astype(np.uint8)  
120 | 
121 | 	return Inorm
122 | 
123 | def new_color_augmentation(patch_np, kdtree, alpha, beta, shift_value=70, threshold=1000):
124 | 	
125 | 	b = False
126 | 	i = 0
127 | 	
128 | 	pipeline_transform_ = A.Compose([
129 | 		A.HueSaturationValue(hue_shift_limit=(-shift_value,shift_value),sat_shift_limit=(-shift_value,shift_value),val_shift_limit=(-shift_value,shift_value),always_apply=True),
130 | 	])
131 | 	
132 | 	while (b==False and i<threshold):
133 | 		
134 | 		A_np = pipeline_transform_(image=patch_np)['image']
135 | 
136 | 		try:
137 | 			HE_ref = H_E_Staining(A_np)
138 | 			point = np.reshape(HE_ref, 6)
139 | 			
140 | 			_, points_indeces = kdtree.query(point, k = alpha, distance_upper_bound = beta)
141 | 			
142 | 			if (unique_elements(points_indeces)):
143 | 				
144 | 				b = True
145 | 				
146 | 			else:
147 | 				
148 | 				i = i + 1
149 | 
150 | 		except:
151 | 			print("H&E not valid")
152 | 	
153 | 
154 | 	return A_np, HE_ref
155 | 
156 | 
157 | def new_stain_augmentation(patch_np, kdtree, alpha, beta, sigma1, sigma2, threshold=1000000):
158 | 	
159 | 	b = False
160 | 	i = 0
161 | 
162 | 	
163 | 	data = kdtree.data
164 | 	
165 | 	while (b==False and i<threshold):
166 | 		
167 | 		idx_HE_ref = np.random.choice(data.shape[0])
168 | 		
169 | 		HE_ref = data[idx_HE_ref]
170 | 		#print(HE_ref)
171 | 		HE_ref = np.reshape(HE_ref, (3,2))
172 | 		
173 | 		alpha_sig = np.random.uniform(1 - sigma1, 1 + sigma1)
174 | 		beta_sig = np.random.uniform(-sigma2, sigma2)
175 | 		
176 | 		HE_ref *= alpha_sig 
177 | 		HE_ref += beta_sig
178 | 		
179 | 		#print(HE_ref)
180 | 		point = np.reshape(HE_ref, 6)
181 | 		
182 | 		_, points_indeces = kdtree.query(point, k = alpha, distance_upper_bound = beta)
183 | 		
184 | 		if (unique_elements(points_indeces)):
185 | 			
186 | 			b = True
187 | 			
188 | 		else:
189 | 			
190 | 			i = i + 1
191 | 	
192 | 	A_np = normalizeStaining(patch_np, HE_ref)
193 | 
194 | 	return A_np, HE_ref
195 | 
196 | 


--------------------------------------------------------------------------------
/scripts/Training_new_augmentation.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | from torch.utils import data
  3 | import numpy as np
  4 | import pandas as pd
  5 | from PIL import Image
  6 | import albumentations as A
  7 | import time
  8 | import torch.nn.functional as F
  9 | import matplotlib.pyplot as plt
 10 | from matplotlib.pyplot import imshow
 11 | import torch.utils.data
 12 | from sklearn import metrics 
 13 | import os
 14 | import shutil
 15 | import sys, getopt
 16 | import warnings 
 17 | 
 18 | import pickle
 19 | import argparse
 20 | from scipy.spatial import KDTree, cKDTree
 21 | from tqdm import tqdm
 22 | 
 23 | from Augmentation import new_stain_augmentation, new_color_augmentation, normalizeStaining, unique_elements, H_E_Staining
 24 | 
 25 | warnings.filterwarnings("ignore")
 26 | 
 27 | argv = sys.argv[1:]
 28 | 
 29 | print("CUDA current device " + str(torch.cuda.current_device()))
 30 | print("CUDA devices available " + str(torch.cuda.device_count()))
 31 | if torch.cuda.is_available():
 32 | 	device = torch.device("cuda")
 33 | 	print("working on gpu")
 34 | else:
 35 | 	device = torch.device("cpu")
 36 | 	print("working on cpu")
 37 | print(torch.backends.cudnn.version())
 38 | 
 39 | #parser parameters
 40 | parser = argparse.ArgumentParser(description='Configurations to train models.')
 41 | parser.add_argument('-n', '--N_EXP', help='number of experiment',type=int, default=0)
 42 | parser.add_argument('-c', '--CNN', help='cnn_to_use',type=str, default='densenet121')
 43 | parser.add_argument('-b', '--BATCH_SIZE', help='batch_size',type=int, default=32)
 44 | parser.add_argument('-e', '--EPOCHS', help='epochs to train',type=int, default=10)
 45 | parser.add_argument('-t', '--TASK', help='task (binary/multilabel)',type=str, default='new_augment')
 46 | parser.add_argument('-f', '--features', help='features_to_use: embedding (True) or features from CNN (False)',type=bool, default=True)
 47 | parser.add_argument('-x', '--extend', help='extend the stored stainings with the trainign data',type=bool, default=False)
 48 | parser.add_argument('-o', '--output', help='output_folder_where_to_store_weights',type=str, default='path_output')
 49 | parser.add_argument('-d', '--database', help='h&e database',type=str, default='path_database')
 50 | parser.add_argument('-i', '--input', help='input csv (path patch, label)',type=str, default='path_input')
 51 | parser.add_argument('-a', '--augmentation', help='type of augmentation: color, stain, he',type=str, default='color')
 52 | 
 53 | 
 54 | args = parser.parse_args()
 55 | 
 56 | N_EXP = args.N_EXP
 57 | N_EXP_str = str(N_EXP)
 58 | CNN_TO_USE = args.CNN
 59 | BATCH_SIZE = args.BATCH_SIZE
 60 | BATCH_SIZE_str = str(BATCH_SIZE)
 61 | EPOCHS = args.EPOCHS
 62 | EPOCHS_str = EPOCHS
 63 | TASK = args.TASK
 64 | EMBEDDING_bool = args.features
 65 | EXTEND = args.extend
 66 | PATH_OUTPUT = args.output
 67 | PATH_DATABASE = args.database
 68 | PATH_INPUT = args.input
 69 | TYPE_AUGMENTATION = args.augmentation
 70 | 
 71 | #np.random.seed(N_EXP)
 72 | 
 73 | def create_dir(directory):
 74 | 	if not os.path.isdir(directory):
 75 | 		try:
 76 | 			os.mkdir(directory)
 77 | 		except OSError:
 78 | 			print ("Creation of the directory %s failed" % directory)
 79 | 		else:
 80 | 			print ("Successfully created the directory %s " % directory) 
 81 | 
 82 | models_path = PATH_OUTPUT
 83 | os.makedirs(models_path,exist_ok=True)
 84 | 
 85 | checkpoint_path = models_path+'checkpoints/'
 86 | os.makedirs(checkpoint_path,exist_ok=True)
 87 | 
 88 | model_path = models_path+'model.pt'
 89 | 
 90 | FOLDER_KD = os.path.split(PATH_DATABASE)[0] + '/'
 91 | 
 92 | fname = PATH_DATABASE
 93 | 
 94 | 
 95 | with open(fname, 'rb') as f:
 96 | 	kdtree = pickle.load(f)
 97 | 
 98 | paths_folder = PATH_INPUT
 99 | 
100 | #import csv
101 | train_csv = paths_folder+'train_patches.csv'
102 | train_dataset = pd.read_csv(train_csv, sep=',',header=None).values
103 | 
104 | valid_csv = paths_folder+'valid_patches.csv'
105 | valid_dataset = pd.read_csv(valid_csv, sep=',',header=None).values
106 | 
107 | imageNet_weights = True
108 | 
109 | 
110 | class ImbalancedDatasetSampler(torch.utils.data.sampler.Sampler):
111 | 	"""Samples elements randomly from a given list of indices for imbalanced dataset
112 | 	Arguments:
113 | 		indices (list, optional): a list of indices
114 | 		num_samples (int, optional): number of samples to draw
115 | 	"""
116 | 
117 | 	def __init__(self, dataset, indices=None, num_samples=None):
118 | 				
119 | 		# if indices is not provided, 
120 | 		# all elements in the dataset will be considered
121 | 		self.indices = list(range(len(dataset)))             if indices is None else indices
122 | 			
123 | 		# if num_samples is not provided, 
124 | 		# draw `len(indices)` samples in each iteration
125 | 		self.num_samples = len(self.indices)             if num_samples is None else num_samples
126 | 			
127 | 		# distribution of classes in the dataset 
128 | 		label_to_count = {}
129 | 		for idx in self.indices:
130 | 			label = self._get_label(dataset, idx)
131 | 			if label in label_to_count:
132 | 				label_to_count[label] += 1
133 | 			else:
134 | 				label_to_count[label] = 1
135 | 				
136 | 		# weight for each sample
137 | 		weights = [1.0 / label_to_count[self._get_label(dataset, idx)]
138 | 				   for idx in self.indices]
139 | 		self.weights = torch.DoubleTensor(weights)
140 | 
141 | 	def _get_label(self, dataset, idx):
142 | 		return dataset[idx,1]
143 | 				
144 | 	def __iter__(self):
145 | 		return (self.indices[i] for i in torch.multinomial(
146 | 			self.weights, self.num_samples, replacement=True))
147 | 
148 | 	def __len__(self):
149 | 		return self.num_samples
150 | 
151 | #MODEL DEFINITION
152 | pre_trained_network = torch.hub.load('pytorch/vision:v0.4.2', CNN_TO_USE, pretrained=imageNet_weights)
153 | 
154 | if (('resnet' in CNN_TO_USE) or ('resnext' in CNN_TO_USE)):
155 | 	fc_input_features = pre_trained_network.fc.in_features
156 | elif (('densenet' in CNN_TO_USE)):
157 | 	fc_input_features = pre_trained_network.classifier.in_features
158 | elif ('mobilenet' in CNN_TO_USE):
159 | 	fc_input_features = pre_trained_network.classifier[1].in_features
160 | 
161 | 
162 | from torch.autograd import Function
163 | class ReverseLayerF(Function):
164 | 
165 | 	@staticmethod
166 | 	def forward(ctx, x, alpha):
167 | 		ctx.alpha = alpha
168 | 
169 | 		return x.view_as(x)
170 | 
171 | 	@staticmethod
172 | 	def backward(ctx, grad_output):
173 | 		output = grad_output.neg() * ctx.alpha
174 | 
175 | 		return output, None
176 | 
177 | class domain_predictor(torch.nn.Module):
178 | 	def __init__(self, n_centers):
179 | 		super(domain_predictor, self).__init__()
180 | 		# domain predictor
181 | 		self.fc_feat_in = fc_input_features
182 | 		self.n_centers = n_centers
183 | 
184 | 		if (EMBEDDING_bool==True):
185 | 			
186 | 			if ('resnet18' in CNN_TO_USE):
187 | 				self.E = 128
188 | 
189 | 			elif ('resnet34' in CNN_TO_USE):
190 | 				self.E = 128
191 | 
192 | 			elif ('resnet50' in CNN_TO_USE):
193 | 				self.E = 256
194 | 			
195 | 			elif ('densenet121' in CNN_TO_USE):
196 | 				self.E = 128
197 | 				
198 | 			
199 | 			self.domain_embedding = torch.nn.Linear(in_features=self.fc_feat_in, out_features=self.E)
200 | 			self.domain_classifier = torch.nn.Linear(in_features=self.E, out_features=self.n_centers)
201 | 
202 | 	def forward(self, x):
203 | 
204 | 		dropout = torch.nn.Dropout(p=0.2)
205 | 		m_binary = torch.nn.Sigmoid()
206 | 
207 | 		domain_emb = self.domain_embedding(x)
208 | 		domain_emb = dropout(domain_emb)
209 | 		domain_prob = self.domain_classifier(domain_emb)
210 | 
211 | 		#domain_prob = m_binary(domain_prob)
212 | 
213 | 		return domain_prob
214 | 
215 | class CNN_model_multitask(torch.nn.Module):
216 | 	def __init__(self):
217 | 		"""
218 | 		In the constructor we instantiate two nn.Linear modules and assign them as
219 | 		member variables.
220 | 		"""
221 | 		super(CNN_model_multitask, self).__init__()
222 | 		self.conv_layers = torch.nn.Sequential(*list(pre_trained_network.children())[:-1])
223 | 
224 | 		if (torch.cuda.device_count()>1):
225 | 			self.conv_layers = torch.nn.DataParallel(self.conv_layers)
226 | 		
227 | 		self.fc_feat_in = fc_input_features
228 | 		self.N_CLASSES = 4
229 | 
230 | 		if (EMBEDDING_bool==True):
231 | 
232 | 			if ('resnet18' in CNN_TO_USE):
233 | 				self.E = 128
234 | 				self.L = self.E
235 | 				self.D = 64
236 | 				self.K = self.N_CLASSES
237 | 
238 | 			elif ('resnet34' in CNN_TO_USE):
239 | 				self.E = 128
240 | 				self.L = self.E
241 | 				self.D = 64
242 | 				self.K = self.N_CLASSES
243 | 				#self.K = 1
244 | 			elif ('resnet50' in CNN_TO_USE):
245 | 				self.E = 256
246 | 				self.L = self.E
247 | 				self.D = 128
248 | 				self.K = self.N_CLASSES
249 | 			elif ('densenet121' in CNN_TO_USE):
250 | 				self.E = 128
251 | 				self.L = self.E
252 | 				self.D = 64
253 | 				self.K = self.N_CLASSES
254 | 
255 | 			#self.embedding = siamese_model.embedding
256 | 			self.embedding = torch.nn.Linear(in_features=self.fc_feat_in, out_features=self.E)
257 | 			self.embedding_fc = torch.nn.Linear(in_features=self.E, out_features=self.N_CLASSES)
258 | 
259 | 		else:
260 | 			self.fc = torch.nn.Linear(in_features=self.fc_feat_in, out_features=self.N_CLASSES)
261 | 			
262 | 			if ('resnet18' in CNN_TO_USE):
263 | 				self.L = fc_input_features
264 | 				self.D = 128
265 | 				self.K = self.N_CLASSES
266 | 
267 | 			elif ('resnet34' in CNN_TO_USE):
268 | 				self.L = fc_input_features
269 | 				self.D = 128
270 | 				self.K = self.N_CLASSES
271 | 
272 | 			elif ('resnet50' in CNN_TO_USE):
273 | 				self.L = self.E
274 | 				self.D = 256
275 | 				self.K = self.N_CLASSES		
276 | 			elif ('densenet121' in CNN_TO_USE):
277 | 				self.E = 128
278 | 				self.L = self.E
279 | 				self.D = 64
280 | 				self.K = self.N_CLASSES
281 | 		
282 | 		self.domain_predictor = domain_predictor(6)
283 | 
284 | 	def forward(self, x, mode, alpha):
285 | 			"""
286 | 			In the forward function we accept a Tensor of input data and we must return
287 | 			a Tensor of output data. We can use Modules defined in the constructor as
288 | 			well as arbitrary operators on Tensors.
289 | 			"""
290 | 			#if used attention pooling
291 | 			A = None
292 | 			#m = torch.nn.Softmax(dim=1)
293 | 			m_binary = torch.nn.Sigmoid()
294 | 			m_multiclass = torch.nn.Softmax()
295 | 			dropout = torch.nn.Dropout(p=0.2)
296 | 			
297 | 			if x is not None:
298 | 				#print(x.shape)
299 | 				conv_layers_out=self.conv_layers(x)
300 | 				#print(x.shape)
301 | 				if ('densenet' in CNN_TO_USE):
302 | 					n = torch.nn.AdaptiveAvgPool2d((1,1))
303 | 					conv_layers_out = n(conv_layers_out)
304 | 				
305 | 				conv_layers_out = conv_layers_out.view(-1, self.fc_feat_in)
306 | 
307 | 			#print(conv_layers_out.shape)
308 | 
309 | 			if ('mobilenet' in CNN_TO_USE):
310 | 				dropout = torch.nn.Dropout(p=0.2)
311 | 				conv_layers_out = dropout(conv_layers_out)
312 | 			#print(conv_layers_out.shape)
313 | 
314 | 			if (EMBEDDING_bool==True):
315 | 				embedding_layer = self.embedding(conv_layers_out)
316 | 				features_to_return = embedding_layer
317 | 
318 | 				embedding_layer = dropout(embedding_layer)
319 | 				logits = self.embedding_fc(embedding_layer)
320 | 
321 | 			else:
322 | 				logits = self.fc(conv_layers_out)
323 | 				features_to_return = conv_layers_out
324 | 
325 | 			output_fcn = m_multiclass(logits)
326 | 
327 | 			if (mode=='train'):
328 | 				reverse_feature = ReverseLayerF.apply(conv_layers_out, alpha)
329 | 
330 | 				output_domain = self.domain_predictor(reverse_feature)
331 | 				output_fcn = m_multiclass(logits)
332 | 
333 | 				return logits, output_fcn, output_domain
334 | 
335 | 			return logits, output_fcn
336 | 
337 | 
338 | class CNN_model(torch.nn.Module):
339 | 	def __init__(self):
340 | 		"""
341 | 		In the constructor we instantiate two nn.Linear modules and assign them as
342 | 		member variables.
343 | 		"""
344 | 		super(CNN_model, self).__init__()
345 | 		self.conv_layers = torch.nn.Sequential(*list(pre_trained_network.children())[:-1])
346 | 
347 | 		if (torch.cuda.device_count()>1):
348 | 			self.conv_layers = torch.nn.DataParallel(self.conv_layers)
349 | 		
350 | 		self.fc_feat_in = fc_input_features
351 | 		self.N_CLASSES = 4
352 | 
353 | 		if (EMBEDDING_bool==True):
354 | 			if ('resnet18' in CNN_TO_USE):
355 | 				self.E = 128
356 | 				self.L = self.E
357 | 				self.D = 64
358 | 				self.K = self.N_CLASSES
359 | 
360 | 			elif ('resnet34' in CNN_TO_USE):
361 | 				self.E = 128
362 | 				self.L = self.E
363 | 				self.D = 64
364 | 				self.K = self.N_CLASSES
365 | 				#self.K = 1
366 | 			elif ('resnet50' in CNN_TO_USE):
367 | 				self.E = 256
368 | 				self.L = self.E
369 | 				self.D = 128
370 | 				self.K = self.N_CLASSES
371 | 			elif ('densenet121' in CNN_TO_USE):
372 | 				self.E = 128
373 | 				self.L = self.E
374 | 				self.D = 64
375 | 				self.K = self.N_CLASSES
376 | 
377 | 			#self.embedding = siamese_model.embedding
378 | 			self.embedding = torch.nn.Linear(in_features=self.fc_feat_in, out_features=self.E)
379 | 			self.embedding_fc = torch.nn.Linear(in_features=self.E, out_features=self.N_CLASSES)
380 | 
381 | 		else:
382 | 			self.fc = torch.nn.Linear(in_features=self.fc_feat_in, out_features=self.N_CLASSES)
383 | 			
384 | 			if ('resnet18' in CNN_TO_USE):
385 | 				self.L = fc_input_features
386 | 				self.D = 128
387 | 				self.K = self.N_CLASSES
388 | 
389 | 			elif ('resnet34' in CNN_TO_USE):
390 | 				self.L = fc_input_features
391 | 				self.D = 128
392 | 				self.K = self.N_CLASSES
393 | 
394 | 			elif ('resnet50' in CNN_TO_USE):
395 | 				self.L = fc_input_features
396 | 				self.D = 256
397 | 				self.K = self.N_CLASSES	
398 | 
399 | 			elif ('densenet121' in CNN_TO_USE):
400 | 				self.L = fc_input_features
401 | 				self.D = 64
402 | 				self.K = self.N_CLASSES
403 | 	
404 | 
405 | 	def forward(self, x, conv_layers_out):
406 | 			"""
407 | 			In the forward function we accept a Tensor of input data and we must return
408 | 			a Tensor of output data. We can use Modules defined in the constructor as
409 | 			well as arbitrary operators on Tensors.
410 | 			"""
411 | 			#if used attention pooling
412 | 			A = None
413 | 			#m = torch.nn.Softmax(dim=1)
414 | 			m_binary = torch.nn.Sigmoid()
415 | 			m_multiclass = torch.nn.Softmax()
416 | 
417 | 			dropout = torch.nn.Dropout(p=0.2)
418 | 			
419 | 			if x is not None:
420 | 				#print(x.shape)
421 | 				conv_layers_out=self.conv_layers(x)
422 | 				#print(x.shape)
423 | 
424 | 				if ('densenet' in CNN_TO_USE):
425 | 					n = torch.nn.AdaptiveAvgPool2d((1,1))
426 | 					conv_layers_out = n(conv_layers_out)
427 | 				
428 | 				conv_layers_out = conv_layers_out.view(-1, self.fc_feat_in)
429 | 
430 | 			#print(conv_layers_out.shape)
431 | 
432 | 			if ('mobilenet' in CNN_TO_USE):
433 | 				dropout = torch.nn.Dropout(p=0.2)
434 | 				conv_layers_out = dropout(conv_layers_out)
435 | 			#print(conv_layers_out.shape)
436 | 
437 | 			if (EMBEDDING_bool==True):
438 | 				embedding_layer = self.embedding(conv_layers_out)
439 | 				features_to_return = embedding_layer
440 | 
441 | 				embedding_layer = dropout(embedding_layer)
442 | 				logits = self.embedding_fc(embedding_layer)
443 | 
444 | 			else:
445 | 				logits = self.fc(conv_layers_out)
446 | 				features_to_return = conv_layers_out
447 | 
448 | 			output_fcn = m_multiclass(logits)
449 | 			
450 | 			return logits, output_fcn 
451 | 
452 | if (TYPE_AUGMENTATION=='he'):
453 | 	model = CNN_model_multitask()
454 | else:
455 | 	model = CNN_model()
456 | 
457 | #DATA AUGMENTATION
458 | from torchvision import transforms
459 | prob = 0.5
460 | 
461 | pipeline_transform = A.Compose([
462 | 		A.VerticalFlip(p=prob),
463 | 		A.HorizontalFlip(p=prob),
464 | 		A.RandomRotate90(p=prob),
465 | 		#A.ElasticTransform(alpha=0.1,p=prob),
466 | 		#A.HueSaturationValue(hue_shift_limit=(-9),sat_shift_limit=25,val_shift_limit=10,p=prob),
467 | 		])
468 | 
469 | #DATA NORMALIZATION
470 | preprocess = transforms.Compose([
471 | 	transforms.ToTensor(),
472 | 	transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
473 | ])
474 | 
475 | def extend_stains(kdtree, new_data, save_new_array = False, PERC=1.0):
476 | 
477 | 	HEs_new_stains = []
478 | 
479 | 	threshold_value = int(len(new_data)*PERC)
480 | 
481 | 	i = 0
482 | 		
483 | 	HEs_general = kdtree.data
484 | 
485 | 	np.random.shuffle(new_data)
486 | 
487 | 	for i in tqdm(range(threshold_value)):
488 | 
489 | 		patch = new_data[i,0]
490 | 		
491 | 		img = Image.open(patch)
492 | 		img_np = np.asarray(img)
493 | 		
494 | 		HE = H_E_Staining(img_np)
495 | 		
496 | 		HE = np.reshape(HE, 6)
497 | 		HEs_new_stains.append(HE)
498 | 		
499 | 		img.close()
500 | 		
501 | 		#i = i + 1
502 | 
503 | 	HEs_new_stains = np.array(HEs_new_stains)
504 | 	HEs_general = np.append(HEs_general,HEs_new_stains,axis=0)
505 | 
506 | 	new_kdtree = cKDTree(HEs_general)
507 | 
508 | 	if (save_new_array==True):
509 | 		print("EXTENDING STAINS")
510 | 		fname = FOLDER_KD + 'kdtree_extended.pickle'
511 | 
512 | 		with open(fname, 'wb') as f:
513 | 			pickle.dump(kdtree, f)
514 | 
515 | 		print("EXTENSION DONE")
516 | 
517 | 	return new_kdtree
518 | 
519 | sigma_perturb = 0.1
520 | nearest_neighbours = 5
521 | 
522 | sigma1 = 0.7
523 | sigma2 = 0.7
524 | 
525 | alpha = nearest_neighbours
526 | beta = sigma_perturb
527 | 
528 | class Dataset_patches(data.Dataset):
529 | 
530 | 	def __init__(self, list_IDs, labels, mode):
531 | 
532 | 		self.labels = labels
533 | 		self.list_IDs = list_IDs
534 | 		self.mode = mode
535 | 		
536 | 	def __len__(self):
537 | 
538 | 		return len(self.list_IDs)
539 | 
540 | 	def __getitem__(self, index):
541 | 
542 | 		# Select sample
543 | 		ID = self.list_IDs[index]
544 | 		# Load data and get label
545 | 		X = Image.open(ID)
546 | 		X = np.asarray(X)
547 | 		y = self.labels[index]
548 | 		#data augmentation
549 | 
550 | 		if (self.mode == 'train'):
551 | 			X = pipeline_transform(image=X)['image']
552 | 
553 | 			rand_val = np.random.rand(1)[0]
554 | 			
555 | 			if (rand_val>prob):
556 | 				
557 | 				if (TYPE_AUGMENTATION=='color'):
558 | 					#print("color")
559 | 					X, _ = new_color_augmentation(X, kdtree, alpha, beta)
560 | 
561 | 				elif (TYPE_AUGMENTATION=='stain'):
562 | 					#print("stain")
563 | 					X, _ = new_stain_augmentation(X,kdtree, alpha, beta, sigma1, sigma2)
564 | 
565 | 				elif (TYPE_AUGMENTATION=='he'):
566 | 					#print("color")
567 | 					X, h_e_matrix  = new_color_augmentation(X, kdtree, alpha, beta)
568 | 
569 | 			if (TYPE_AUGMENTATION=='he'):
570 | 
571 | 				h_e_matrix = H_E_Staining(X)
572 | 				
573 | 				h_e_matrix = np.reshape(h_e_matrix, 6)
574 | 				h_e_matrix = np.asarray(h_e_matrix)
575 | 			else:
576 | 				h_e_matrix = np.asarray([0])
577 | 
578 | 		new_image = np.asarray(X)
579 | 		#data transformation
580 | 		input_tensor = preprocess(new_image)
581 | 				
582 | 		return input_tensor, np.asarray(y), h_e_matrix
583 | 
584 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
585 | #device = torch.device('cpu')
586 | 
587 | # Parameters
588 | 
589 | num_workers = 2
590 | params_train = {'batch_size': BATCH_SIZE,
591 | 		  #'shuffle': True,
592 | 		  'sampler': ImbalancedDatasetSampler(train_dataset),
593 | 		  'num_workers': num_workers}
594 | 
595 | params_valid = {'batch_size': BATCH_SIZE,
596 | 		  'shuffle': True,
597 | 		  #'sampler': ImbalancedDatasetSampler(valid_dataset),
598 | 		  'num_workers': num_workers}
599 | 
600 | params_test = {'batch_size': BATCH_SIZE,
601 | 		  'shuffle': True,
602 | 		  #'sampler': ImbalancedDatasetSampler(test_dataset),
603 | 		  'num_workers': num_workers}
604 | 
605 | max_epochs = int(EPOCHS_str)
606 | 
607 | 
608 | 
609 | # In[28]:
610 | 
611 | 
612 | #CREATE GENERATORS
613 | #train
614 | training_set = Dataset_patches(train_dataset[:,0], train_dataset[:,1],'train')
615 | training_generator = data.DataLoader(training_set, **params_train)
616 | 
617 | validation_set = Dataset_patches(valid_dataset[:,0], valid_dataset[:,1],'valid')
618 | validation_generator = data.DataLoader(validation_set, **params_valid)
619 | 
620 | 
621 | #semi-weakly supervision
622 | 
623 | # Find total parameters and trainable parameters
624 | total_params = sum(p.numel() for p in model.parameters())
625 | print(f'{total_params:,} total parameters.')
626 | total_trainable_params = sum(
627 | 	p.numel() for p in model.parameters() if p.requires_grad)
628 | print(f'{total_trainable_params:,} training parameters.')
629 | 
630 | class_sample_count = np.unique(train_dataset[:,1], return_counts=True)[1]
631 | weight = class_sample_count / len(train_dataset[:,1])
632 | #for avoiding propagation of fake benign class
633 | samples_weight = torch.from_numpy(weight).type(torch.FloatTensor)
634 | 
635 | class RMSELoss(torch.nn.Module):
636 |     def __init__(self, eps=1e-6):
637 |         super().__init__()
638 |         self.mse = torch.nn.MSELoss()
639 |         self.eps = eps
640 |         
641 |     def forward(self,yhat,y):
642 |         loss = torch.sqrt(self.mse(yhat,y) + self.eps)
643 |         return loss
644 | 
645 | import torch.optim as optim
646 | 
647 | criterion_domain = RMSELoss()
648 | criterion = torch.nn.CrossEntropyLoss()
649 | 
650 | num_epochs = EPOCHS
651 | epoch = 0
652 | early_stop_cont = 0
653 | EARLY_STOP_NUM = 5
654 | #weight_decay = 1e-4
655 | weight_decay = 0
656 | lr = 1e-3
657 | 
658 | optimizer = optim.Adam(model.parameters(),lr=lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=weight_decay, amsgrad=True)
659 | model.to(device)
660 | 
661 | if (EXTEND==True):
662 | 
663 | 	new_kdtree = extend_stains(kdtree, train_dataset, save_new_array = False, PERC=0.5)
664 | 	kdtree = new_kdtree
665 | 
666 | else:
667 | 
668 | 	#fname = FOLDER_KD + 'kdtree_TCGA_ExaMode_extended_prostate.pickle'
669 | 	fname = FOLDER_KD + 'kdtree_TCGA_ExaMode.pickle'
670 | 
671 | 	with open(fname, 'rb') as f:
672 | 		kdtree = pickle.load(f)
673 | 
674 | 
675 | def evaluate_validation_set(generator):
676 | 	#accumulator for validation set
677 | 	y_pred = []
678 | 	y_true = []
679 | 
680 | 	valid_loss = 0.0
681 | 
682 | 	with torch.no_grad():
683 | 		j = 0
684 | 		for inputs,labels, _ in generator:
685 | 			inputs, labels = inputs.to(device), labels.to(device)
686 | 
687 | 			# forward + backward + optimize
688 | 			logits, outputs = model(inputs, None)
689 | 
690 | 			loss = criterion(logits, labels)
691 | 			#outputs = F.softmax(outputs)
692 | 
693 | 			valid_loss = valid_loss + ((1 / (j+1)) * (loss.item() - valid_loss)) 
694 | 			
695 | 			outputs_np = outputs.cpu().data.numpy()
696 | 			labels_np = labels.cpu().data.numpy()
697 | 			outputs_np = np.argmax(outputs_np, axis=1)
698 | 
699 | 			y_true = np.append(y_true, outputs_np)
700 | 			y_pred = np.append(y_pred, labels_np)
701 | 
702 | 			j = j+1			
703 | 
704 | 		acc_valid = metrics.accuracy_score(y_true=y_true, y_pred=y_pred)
705 | 		kappa_valid = metrics.cohen_kappa_score(y1=y_true,y2=y_pred, weights='quadratic')
706 | 		print("loss: " + str(valid_loss) + ", accuracy: " + str(acc_valid) + ", kappa score: " + str(kappa_valid))
707 | 		
708 | 	return valid_loss
709 | # In[35]:
710 | 
711 | best_loss_valid = 100000.0
712 | 
713 | losses_train = []
714 | losses_valid = []
715 | 
716 | 
717 | lambda_val = 0.5 
718 | 
719 | while (epoch<num_epochs and early_stop_cont<EARLY_STOP_NUM):
720 | 	
721 | 	y_true = []
722 | 	y_pred = []
723 | 
724 | 	#loss functions outputs and network
725 | 	train_loss = 0.0
726 | 
727 | 	train_loss_patches = 0.0
728 | 	train_loss_domain = 0.0
729 | 
730 | 	is_best = False
731 | 	
732 | 	i = 0
733 | 	
734 | 	model.train()
735 | 
736 | 	tot_iterations = int(len(train_dataset)/BATCH_SIZE)
737 | 	
738 | 	for inputs,labels, h_e_matrices in training_generator:
739 | 		inputs, labels = inputs.to(device), labels.to(device)
740 | 		h_e_matrices = h_e_matrices.type(torch.FloatTensor).to(device)
741 | 
742 | 
743 | 		if (TYPE_AUGMENTATION=='he'):
744 | 			p = float(i + epoch * tot_iterations) / num_epochs / tot_iterations
745 | 
746 | 			alpha = 2. / (1. + np.exp(-10 * p)) - 1
747 | 
748 | 		# zero the parameter gradients
749 | 		optimizer.zero_grad()
750 | 		
751 | 		# forward + backward + optimize
752 | 		if (TYPE_AUGMENTATION=='he'):
753 | 			logits, outputs, pred_domain = model(inputs, 'train', alpha)
754 | 			#pred_domain = pred_domain.view(-1)
755 | 
756 | 			loss_patches = criterion(logits, labels)
757 | 
758 | 			loss_domains = lambda_val * criterion_domain(pred_domain, h_e_matrices)
759 | 
760 | 			loss = loss_patches + loss_domains
761 | 
762 | 			loss.backward()
763 | 			optimizer.step()
764 | 			
765 | 			train_loss = train_loss + ((1 / (i+1)) * (loss.item() - train_loss))   
766 | 			train_loss_patches = train_loss_patches + ((1 / (i+1)) * (loss_patches.item() - train_loss_patches)) 
767 | 			train_loss_domain = train_loss_domain + ((1 / (i+1)) * (loss_domains.item() - train_loss_domain)) 
768 | 
769 | 			#outputs = F.softmax(outputs)
770 | 			#accumulate values
771 | 			outputs_np = outputs.cpu().data.numpy()
772 | 			labels_np = labels.cpu().data.numpy()
773 | 			outputs_np = np.argmax(outputs_np, axis=1)
774 | 
775 | 			y_true = np.append(y_true, outputs_np)
776 | 			y_pred = np.append(y_pred, labels_np)
777 | 
778 | 			i = i+1
779 | 
780 | 			if (i%100==0):
781 | 
782 | 				print("loss: " + str(train_loss) + " loss patches: " + str(train_loss_patches) + " loss domains: " + str(train_loss_domain))
783 | 
784 | 				print("["+str(i)+"/"+str(tot_iterations)+"]")
785 | 				acc = metrics.accuracy_score(y_true=y_true, y_pred=y_pred)
786 | 				kappa = metrics.cohen_kappa_score(y1=y_true,y2=y_pred, weights='quadratic')
787 | 
788 | 				print("accuracy: " + str(acc))
789 | 				print("kappa score: " + str(kappa))
790 | 		
791 | 		else:
792 | 			logits, outputs = model(inputs, None)
793 | 			#print(logits.shape,labels.shape)
794 | 			loss = criterion(logits, labels)
795 | 
796 | 			loss.backward()
797 | 			optimizer.step()
798 | 			
799 | 			train_loss = train_loss + ((1 / (i+1)) * (loss.item() - train_loss))   
800 | 			#outputs = F.softmax(outputs)
801 | 			#accumulate values
802 | 			outputs_np = outputs.cpu().data.numpy()
803 | 			labels_np = labels.cpu().data.numpy()
804 | 			outputs_np = np.argmax(outputs_np, axis=1)
805 | 
806 | 			y_true = np.append(y_true, outputs_np)
807 | 			y_pred = np.append(y_pred, labels_np)
808 | 
809 | 			i = i+1
810 | 			
811 | 			if (i%100==0):
812 | 				print("["+str(i)+"/"+str(tot_iterations)+"]")
813 | 				acc = metrics.accuracy_score(y_true=y_true, y_pred=y_pred)
814 | 				kappa = metrics.cohen_kappa_score(y1=y_true,y2=y_pred, weights='quadratic')
815 | 
816 | 				print("accuracy: " + str(acc))
817 | 				print("kappa score: " + str(kappa))
818 | 
819 | 
820 | 		
821 | 		optimizer.zero_grad()
822 | 		torch.cuda.empty_cache()
823 | 
824 | 	model.eval()
825 | 
826 | 	print("epoch "+str(epoch)+ " train loss: " + str(train_loss) + " acc_train: " + str(acc))
827 | 	
828 | 	print("evaluating validation")
829 | 	valid_loss = evaluate_validation_set(validation_generator)
830 | 	
831 | 	if (best_loss_valid>valid_loss):
832 | 		print ("=> Saving a new best model")
833 | 		print("previous loss TMA: " + str(best_loss_valid) + ", new loss function TMA: " + str(valid_loss))
834 | 		best_loss_valid = valid_loss
835 | 		torch.save(model, model_path)
836 | 		early_stop_cont = 0
837 | 	else:
838 | 		early_stop_cont = early_stop_cont+1
839 | 		
840 | 	epoch = epoch + 1
841 | 	
842 | print('Finished Training')
843 | 


--------------------------------------------------------------------------------