├── .gitattributes ├── .gitignore ├── LICENSE ├── README.md ├── anno └── B_anno.csv ├── datasets ├── __init__.py ├── basic_dataset.py ├── data_folder.py ├── dataloader_prefetch.py ├── separate_files_dataset.py └── single_file_dataset.py ├── documents ├── sample_id_all.tsv ├── sample_id_testset.tsv └── supplementary.pdf ├── environment.yml ├── models ├── __init__.py ├── basic_model.py ├── c_gan_model.py ├── c_gan_ved_model.py ├── losses.py └── networks.py ├── params ├── __init__.py ├── basic_params.py ├── test_params.py ├── train_params.py └── train_test_param.py ├── requirements.txt ├── test.py ├── train.py ├── train_test.py └── util ├── __init__.py ├── preprocess.py ├── util.py └── visualizer.py /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .ipynb_checkpoints 2 | .idea/ 3 | .DS_Store 4 | checkpoints/ 5 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Xiaoyu Zhang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # OmiTrans 2 | ***Please also have a look at our multi-omics multi-task DL freamwork 👀:*** 3 | [OmiEmbed](https://github.com/zhangxiaoyu11/OmiEmbed) 4 | 5 | [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.5728496.svg)](https://doi.org/10.5281/zenodo.5728496) 6 | [![Codacy Badge](https://app.codacy.com/project/badge/Grade/ff39650740bb4973b211a4fcfb6c1695)](https://www.codacy.com/gh/zhangxiaoyu11/OmiTrans/dashboard?utm_source=github.com&utm_medium=referral&utm_content=zhangxiaoyu11/OmiTrans&utm_campaign=Badge_Grade) 7 | [![GitHub license](https://img.shields.io/github/license/Naereen/StrapDown.js.svg)](https://github.com/zhangxiaoyu11/OmiEmbed/blob/main/LICENSE) 8 | ![Safe](https://img.shields.io/badge/Stay-Safe-red?logo=data:image/svg%2bxml;base64,PHN2ZyBpZD0iTGF5ZXJfMSIgZW5hYmxlLWJhY2tncm91bmQ9Im5ldyAwIDAgNTEwIDUxMCIgaGVpZ2h0PSI1MTIiIHZpZXdCb3g9IjAgMCA1MTAgNTEwIiB3aWR0aD0iNTEyIiB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciPjxnPjxnPjxwYXRoIGQ9Im0xNzQuNjEgMzAwYy0yMC41OCAwLTQwLjU2IDYuOTUtNTYuNjkgMTkuNzJsLTExMC4wOSA4NS43OTd2MTA0LjQ4M2g1My41MjlsNzYuNDcxLTY1aDEyNi44MnYtMTQ1eiIgZmlsbD0iI2ZmZGRjZSIvPjwvZz48cGF0aCBkPSJtNTAyLjE3IDI4NC43MmMwIDguOTUtMy42IDE3Ljg5LTEwLjc4IDI0LjQ2bC0xNDguNTYgMTM1LjgyaC03OC4xOHYtODVoNjguMThsMTE0LjM0LTEwMC4yMWMxMi44Mi0xMS4yMyAzMi4wNi0xMC45MiA0NC41LjczIDcgNi41NSAxMC41IDE1LjM4IDEwLjUgMjQuMnoiIGZpbGw9IiNmZmNjYmQiLz48cGF0aCBkPSJtMzMyLjgzIDM0OS42M3YxMC4zN2gtNjguMTh2LTYwaDE4LjU1YzI3LjQxIDAgNDkuNjMgMjIuMjIgNDkuNjMgNDkuNjN6IiBmaWxsPSIjZmZjY2JkIi8+PHBhdGggZD0ibTM5OS44IDc3LjN2OC4wMWMwIDIwLjY1LTguMDQgNDAuMDctMjIuNjQgNTQuNjdsLTExMi41MSAxMTIuNTF2LTIyNi42NmwzLjE4LTMuMTljMTQuNi0xNC42IDM0LjAyLTIyLjY0IDU0LjY3LTIyLjY0IDQyLjYyIDAgNzcuMyAzNC42OCA3Ny4zIDc3LjN6IiBmaWxsPSIjZDAwMDUwIi8+PHBhdGggZD0ibTI2NC42NSAyNS44M3YyMjYuNjZsLTExMi41MS0xMTIuNTFjLTE0LjYtMTQuNi0yMi42NC0zNC4wMi0yMi42NC01NC42N3YtOC4wMWMwLTQyLjYyIDM0LjY4LTc3LjMgNzcuMy03Ny4zIDIwLjY1IDAgNDAuMDYgOC4wNCA1NC42NiAyMi42NHoiIGZpbGw9IiNmZjRhNGEiLz48cGF0aCBkPSJtMjEyLjgzIDM2MC4xMnYzMGg1MS44MnYtMzB6IiBmaWxsPSIjZmZjY2JkIi8+PHBhdGggZD0ibTI2NC42NSAzNjAuMTJ2MzBoMzYuMTRsMzIuMDQtMzB6IiBmaWxsPSIjZmZiZGE5Ii8+PC9nPjwvc3ZnPg==) 9 | [![GitHub Repo stars](https://img.shields.io/github/stars/zhangxiaoyu11/OmiTrans?style=social)](https://github.com/zhangxiaoyu11/OmiTrans/stargazers) 10 | [![GitHub forks](https://img.shields.io/github/forks/zhangxiaoyu11/OmiTrans?style=social)](https://github.com/zhangxiaoyu11/OmiTrans/network/members) 11 | 12 | **The FIRST GANs-based omics-to-omics translation framework** 13 | 14 | **Xiaoyu Zhang** (x.zhang18@imperial.ac.uk) 15 | 16 | Data Science Institute, Imperial College London 17 | 18 | ## Introduction 19 | 20 | OmiTrans is a generative adversarial networks (GANs) based omics-to-omics translation framework. 21 | 22 | ## Getting Started 23 | 24 | ### Prerequisites 25 | - CPU or NVIDIA GPU + CUDA CuDNN 26 | - [Python](https://www.python.org/downloads) 3.6+ 27 | - Python Package Manager 28 | - [Anaconda](https://docs.anaconda.com/anaconda/install) 3 (recommended) 29 | - or [pip](https://pip.pypa.io/en/stable/installing/) 21.0+ 30 | - Python Packages 31 | - [PyTorch](https://pytorch.org/get-started/locally) 1.2+ 32 | - TensorBoard 1.10+ 33 | - Tables 3.6+ 34 | - prefetch-generator 1.0+ 35 | - [Git](https://git-scm.com/book/en/v2/Getting-Started-Installing-Git) 2.7+ 36 | 37 | ### Installation 38 | - Clone the repo 39 | ```bash 40 | git clone https://github.com/zhangxiaoyu11/OmiTrans.git 41 | cd OmiTrans 42 | ``` 43 | - Install the dependencies 44 | - For conda users 45 | ```bash 46 | conda env create -f environment.yml 47 | conda activate omitrans 48 | ``` 49 | - For pip users 50 | ```bash 51 | pip install -r requirements.txt 52 | ``` 53 | 54 | ### Try it out 55 | - Put the gene expression data (A.tsv) and DNA methylation data (B.tsv) in the default data path (./data) 56 | - Train and test using the default settings 57 | ```bash 58 | python train_test.py 59 | ``` 60 | - Check the output files 61 | ```bash 62 | cd checkpoints/test/ 63 | ``` 64 | - Visualise the metrics and losses 65 | ```bash 66 | tensorboard --logdir=tb_log --bind_all 67 | ``` 68 | 69 | ## OmiEmbed 70 | ***Please also have a look at our multi-omics multi-task DL freamwork 👀:*** 71 | [OmiEmbed](https://github.com/zhangxiaoyu11/OmiEmbed) 72 | 73 | ## License 74 | This source code is licensed under the [MIT](https://github.com/zhangxiaoyu11/OmiTrans/blob/master/LICENSE) license. 75 | -------------------------------------------------------------------------------- /datasets/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | This package about data loading and data preprocessing 3 | """ 4 | import os 5 | import numpy as np 6 | import pandas as pd 7 | import importlib 8 | from torch.utils.data import Subset 9 | from datasets.basic_dataset import BasicDataset 10 | from datasets.dataloader_prefetch import DataLoaderPrefetch 11 | from sklearn.model_selection import train_test_split 12 | 13 | 14 | def find_dataset_using_name(dataset_mode): 15 | """ 16 | Get the dataset of certain mode 17 | """ 18 | dataset_filename = "datasets." + dataset_mode + "_dataset" 19 | datasetlib = importlib.import_module(dataset_filename) 20 | 21 | # Instantiate the dataset class 22 | dataset = None 23 | # Change the name format to corresponding class name 24 | target_dataset_name = dataset_mode.replace('_', '') + 'dataset' 25 | for name, cls in datasetlib.__dict__.items(): 26 | if name.lower() == target_dataset_name.lower() \ 27 | and issubclass(cls, BasicDataset): 28 | dataset = cls 29 | 30 | if dataset is None: 31 | raise NotImplementedError("In %s.py, there should be a subclass of BasicDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name)) 32 | 33 | return dataset 34 | 35 | 36 | def create_dataset(param): 37 | """ 38 | Create a dataset given the parameters. 39 | """ 40 | dataset_class = find_dataset_using_name(param.dataset_mode) 41 | # Get an instance of this dataset class 42 | dataset = dataset_class(param) 43 | print("Dataset [%s] was created" % type(dataset).__name__) 44 | 45 | return dataset 46 | 47 | 48 | class CustomDataLoader: 49 | """ 50 | Create a dataloader for certain dataset. 51 | """ 52 | def __init__(self, dataset, param, shuffle=True, enable_drop_last=False): 53 | self.dataset = dataset 54 | self.param = param 55 | 56 | drop_last = False 57 | if enable_drop_last: 58 | if len(dataset) % param.batch_size < 3 * len(param.gpu_ids): 59 | drop_last = True 60 | 61 | # Create dataloader for this dataset 62 | self.dataloader = DataLoaderPrefetch( 63 | dataset, 64 | batch_size=param.batch_size, 65 | shuffle=shuffle, 66 | num_workers=int(param.num_threads), 67 | drop_last=drop_last, 68 | pin_memory=param.set_pin_memory 69 | ) 70 | 71 | def __len__(self): 72 | """Return the number of data in the dataset""" 73 | return len(self.dataset) 74 | 75 | def __iter__(self): 76 | """Return a batch of data""" 77 | for i, data in enumerate(self.dataloader): 78 | yield data 79 | 80 | def get_A_dim(self): 81 | """Return the dimension of first input omics data type""" 82 | return self.dataset.A_dim 83 | 84 | def get_B_dim(self): 85 | """Return the dimension of second input omics data type""" 86 | return self.dataset.B_dim 87 | 88 | def get_omics_dims(self): 89 | """Return a list of omics dimensions""" 90 | return self.dataset.omics_dims 91 | 92 | def get_sample_list(self): 93 | """Return the sample list of the dataset""" 94 | return self.dataset.sample_list 95 | 96 | def get_feature_list_A(self): 97 | """Return the feature list A of the dataset""" 98 | return self.dataset.feature_list_A 99 | 100 | def get_values_max(self): 101 | """Return the maximum target value of the dataset""" 102 | return self.dataset.target_max 103 | 104 | def get_values_min(self): 105 | """Return the minimum target value of the dataset""" 106 | return self.dataset.target_min 107 | 108 | 109 | def create_single_dataloader(param, shuffle=True, enable_drop_last=False): 110 | """ 111 | Create a single dataloader 112 | """ 113 | param.stratify = False 114 | dataset = create_dataset(param) 115 | dataloader = CustomDataLoader(dataset, param, shuffle=shuffle, enable_drop_last=enable_drop_last) 116 | sample_list = dataset.sample_list 117 | 118 | return dataloader, sample_list 119 | 120 | 121 | def create_separate_dataloader(param): 122 | """ 123 | Create set of dataloader (train, val, test). 124 | """ 125 | full_dataset = create_dataset(param) 126 | full_size = len(full_dataset) 127 | full_idx = np.arange(full_size) 128 | 129 | if param.stratify: 130 | train_idx, test_idx = train_test_split(full_idx, 131 | test_size=param.test_ratio, 132 | train_size=param.train_ratio, 133 | shuffle=True, 134 | stratify=full_dataset.labels_array) 135 | else: 136 | train_idx, test_idx = train_test_split(full_idx, 137 | test_size=param.test_ratio, 138 | train_size=param.train_ratio, 139 | shuffle=True) 140 | 141 | val_idx = list(set(full_idx) - set(train_idx) - set(test_idx)) 142 | 143 | train_dataset = Subset(full_dataset, train_idx) 144 | val_dataset = Subset(full_dataset, val_idx) 145 | test_dataset = Subset(full_dataset, test_idx) 146 | 147 | full_dataloader = CustomDataLoader(full_dataset, param) 148 | train_dataloader = CustomDataLoader(train_dataset, param, enable_drop_last=True) 149 | val_dataloader = CustomDataLoader(val_dataset, param, shuffle=False) 150 | test_dataloader = CustomDataLoader(test_dataset, param, shuffle=False) 151 | 152 | return full_dataloader, train_dataloader, val_dataloader, test_dataloader 153 | 154 | 155 | def load_file(param, file_name): 156 | """ 157 | Load data according to the format. 158 | """ 159 | if param.file_format == 'tsv': 160 | file_path = os.path.join(param.data_root, file_name + '.tsv') 161 | print('Loading data from ' + file_path) 162 | df = pd.read_csv(file_path, sep='\t', header=0, index_col=0, na_filter=param.detect_na) 163 | elif param.file_format == 'csv': 164 | file_path = os.path.join(param.data_root, file_name + '.csv') 165 | print('Loading data from ' + file_path) 166 | df = pd.read_csv(file_path, header=0, index_col=0, na_filter=param.detect_na) 167 | elif param.file_format == 'hdf': 168 | file_path = os.path.join(param.data_root, file_name + '.h5') 169 | print('Loading data from ' + file_path) 170 | df = pd.read_hdf(file_path, header=0, index_col=0) 171 | else: 172 | raise NotImplementedError('File format %s is supported' % param.file_format) 173 | return df 174 | 175 | 176 | def save_file(param, dataframe, file_name): 177 | """ 178 | Save the dataframe to disk according to the format 179 | """ 180 | if param.file_format == 'tsv': 181 | output_path = os.path.join(param.checkpoints_dir, param.experiment_name, file_name + '.tsv') 182 | dataframe.to_csv(output_path, sep='\t') 183 | elif param.file_format == 'csv': 184 | output_path = os.path.join(param.checkpoints_dir, param.experiment_name, file_name + '.csv') 185 | dataframe.to_csv(output_path) 186 | elif param.file_format == 'hdf': 187 | output_path = os.path.join(param.checkpoints_dir, param.experiment_name, file_name + '.h5') 188 | dataframe.to_hdf(output_path, key='df', mode='w') 189 | else: 190 | raise NotImplementedError('File format %s is supported' % param.file_format) 191 | -------------------------------------------------------------------------------- /datasets/basic_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | This module implements an abstract base class for datasets. Other datasets can be created from this base class. 3 | """ 4 | import torch.utils.data as data 5 | from abc import ABC, abstractmethod 6 | 7 | 8 | class BasicDataset(data.Dataset, ABC): 9 | """This class is an abstract base class for datasets. 10 | To create a subclass, you need to implement the following three functions: 11 | -- <__init__>: initialize the class, first call BaseDataset.__init__(self, param). 12 | -- <__len__>: return the size of dataset. 13 | -- <__getitem__>: get a data point. 14 | """ 15 | 16 | def __init__(self, param): 17 | """ 18 | Initialize the class, save the parameters in the class 19 | """ 20 | self.param = param 21 | self.sample_list = None 22 | 23 | @abstractmethod 24 | def __len__(self): 25 | """Return the total number of samples in the dataset.""" 26 | return 0 27 | 28 | @abstractmethod 29 | def __getitem__(self, index): 30 | """ 31 | Return a data point and its metadata information. 32 | Parameters: 33 | index - - a integer for data indexing 34 | Returns: 35 | a dictionary of data with their names. It usually contains the data itself and its metadata information. 36 | """ 37 | pass 38 | -------------------------------------------------------------------------------- /datasets/data_folder.py: -------------------------------------------------------------------------------- 1 | """ 2 | Data folder class 3 | """ 4 | import os 5 | import os.path 6 | 7 | # For now only support preprocessed .tsv file with omics matrix 8 | OMICS_EXTENSIONS = [ 9 | '.tsv', 10 | '.csv', 11 | '.h5' 12 | ] 13 | 14 | 15 | def is_omics_file(filename): 16 | return any(filename.endswith(extension) for extension in OMICS_EXTENSIONS) 17 | 18 | 19 | def make_dataset(directory): 20 | """ 21 | Return paths of omics files in the dataset 22 | """ 23 | data_paths = [] 24 | assert os.path.isdir(directory), '%s is not a valid directory!' % directory 25 | 26 | for root, _, file_names in sorted(os.walk(directory)): 27 | for file_name in file_names: 28 | if is_omics_file(file_name): 29 | file_path = os.path.join(root, file_name) 30 | data_paths.append(file_path) 31 | return data_paths 32 | -------------------------------------------------------------------------------- /datasets/dataloader_prefetch.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import DataLoader 2 | from prefetch_generator import BackgroundGenerator 3 | 4 | 5 | class DataLoaderPrefetch(DataLoader): 6 | def __iter__(self): 7 | return BackgroundGenerator(super().__iter__()) 8 | -------------------------------------------------------------------------------- /datasets/separate_files_dataset.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | from datasets.basic_dataset import BasicDataset 3 | from datasets.data_folder import make_dataset 4 | import numpy as np 5 | import pandas as pd 6 | import torch 7 | from natsort import natsorted 8 | 9 | 10 | class SeparateFilesDataset(BasicDataset): 11 | """ 12 | A dataset class for paired omics dataset with separate files. 13 | The data should be prepared in a directory '/path/to/data/'. 14 | The SeparateFilesDataset doesn't support use_sample_list, use_feature_lists and ch_separate. 15 | """ 16 | 17 | def __init__(self, param): 18 | """ 19 | Initialize this dataset class. 20 | """ 21 | BasicDataset.__init__(self, param) 22 | dir_data = os.path.join(param.data_root, 'all') 23 | self.data_paths = natsorted(make_dataset(dir_data)) # get omics data paths 24 | self.omics_dims = [] 25 | A_array_example = np.loadtxt(self.data_paths[0]) 26 | B_array_example = np.loadtxt(self.data_paths[1]) 27 | self.A_dim = len(A_array_example) 28 | self.omics_dims.append(self.A_dim) 29 | self.B_dim = len(B_array_example) 30 | self.omics_dims.append(self.B_dim) 31 | # Get sample list 32 | sample_list_path = os.path.join(param.data_root, 'sample_list.tsv') # get the path of sample list 33 | self.sample_list = np.loadtxt(sample_list_path, delimiter='\t', dtype=' of the model class.""" 33 | model_class = find_model_using_name(model_name) 34 | return model_class.modify_commandline_parameters 35 | 36 | 37 | def create_model(param): 38 | """ 39 | Create a model given the parameters 40 | """ 41 | model = find_model_using_name(param.model) 42 | # Initialize the model 43 | instance = model(param) 44 | print('Model [%s] was created' % type(instance).__name__) 45 | return instance 46 | -------------------------------------------------------------------------------- /models/basic_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from abc import ABC, abstractmethod 4 | from collections import OrderedDict 5 | from . import networks 6 | import numpy as np 7 | import sklearn.metrics as metrics 8 | 9 | 10 | class BasicModel(ABC): 11 | """ 12 | This class is an abstract base class for models. 13 | To create a subclass, you need to implement the following five functions: 14 | -- <__init__>: Initialize the class, first call BaseModel.__init__(self, param) 15 | -- : Add model-specific parameters, and rewrite default values for existing parameters 16 | -- : Unpack input data from the output dictionary of the dataloader 17 | -- : Get the fake omics data 18 | -- : Calculate losses, gradients and update network parameters 19 | """ 20 | 21 | def __init__(self, param): 22 | """ 23 | Initialize the BaseModel class 24 | """ 25 | self.param = param 26 | self.gpu_ids = param.gpu_ids 27 | self.isTrain = param.isTrain 28 | self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') # get device name: CPU or GPU 29 | self.save_dir = os.path.join(param.checkpoints_dir, param.experiment_name) # save all the checkpoints to save_dir, and this is where to load the models 30 | self.load_net_dir = os.path.join(param.checkpoints_dir, param.experiment_to_load) # load pretrained networks from certain experiment folder 31 | 32 | # Improve the performance if the dimensionality and shape of the input data keep the same 33 | torch.backends.cudnn.benchmark = True 34 | 35 | self.loss_names = [] 36 | self.model_names = [] 37 | self.optimizers = [] 38 | self.schedulers = [] 39 | 40 | self.real_A_tensor = [] 41 | self.real_B_tensor = [] 42 | self.fake_A_tensor = [] 43 | 44 | self.data_index = [] # The indexes of input data 45 | 46 | # specify the metrics you want to print out. 47 | self.metric_names = ['MSE', 'RMSE', 'MAE', 'MEDAE', 'R2'] 48 | self.metric_MSE = [] 49 | self.metric_RMSE = [] 50 | self.metric_MAE = [] 51 | self.metric_MEDAE = [] 52 | self.metric_R2 = [] 53 | 54 | self.metric = 0 # used for learning rate policy 'plateau' 55 | 56 | @staticmethod 57 | def modify_commandline_parameters(parser, is_train): 58 | """ 59 | Add model-specific parameters, and rewrite default values for existing parameters. 60 | 61 | Parameters: 62 | parser -- original parameter parser 63 | is_train (bool) -- whether it is currently training phase or test phase. Use this flag to add or change training-specific or test-specific parameters. 64 | 65 | Returns: 66 | The modified parser. 67 | """ 68 | return parser 69 | 70 | @abstractmethod 71 | def set_input(self, input_dict): 72 | """ 73 | Unpack input data from the output dictionary of the dataloader 74 | 75 | Parameters: 76 | input_dict (dict): include the data tensor and its path. This dict looks like this {'A_tensor': A_tensor, 'B_tensor': B_tensor, 'A_path': A_path, 'B_path': B_path} 77 | """ 78 | pass 79 | 80 | @abstractmethod 81 | def forward(self): 82 | """ 83 | Run forward pass 84 | """ 85 | pass 86 | 87 | @abstractmethod 88 | def update(self): 89 | """ 90 | Calculate losses, gradients and update network weights; called in every training iteration 91 | """ 92 | pass 93 | 94 | def set_requires_grad(self, nets, requires_grad=False): 95 | """ 96 | Set requies_grad=Fasle to avoid back propagation 97 | Parameters: 98 | nets (network list) -- a list of networks 99 | requires_grad (bool) -- whether the networks require gradients or not 100 | """ 101 | if not isinstance(nets, list): 102 | nets = [nets] 103 | for net in nets: 104 | if net is not None: 105 | for net_param in net.parameters(): 106 | net_param.requires_grad = requires_grad 107 | 108 | def setup(self, param): 109 | """ 110 | Load and print networks, create schedulers 111 | """ 112 | if self.isTrain: 113 | self.print_networks(param) 114 | # For every optimizer we have a scheduler 115 | self.schedulers = [networks.get_scheduler(optimizer, param) for optimizer in self.optimizers] 116 | 117 | # Loading the networks 118 | if not self.isTrain or param.continue_train: 119 | self.load_networks(param.epoch_to_load) 120 | 121 | def save_networks(self, epoch): 122 | """ 123 | Save all the networks to the disk. 124 | 125 | Parameters: 126 | epoch (str) -- current epoch 127 | """ 128 | for name in self.model_names: 129 | if isinstance(name, str): 130 | save_filename = '{:s}_net_{:s}.pth'.format(epoch, name) 131 | save_path = os.path.join(self.save_dir, save_filename) 132 | # Use the str to get the attribute aka the network (self.netG / self.netD) 133 | net = getattr(self, 'net' + name) 134 | # If we use multi GPUs and apply the data parallel 135 | if len(self.gpu_ids) > 0 and torch.cuda.is_available(): 136 | torch.save(net.module.cpu().state_dict(), save_path) 137 | net.cuda(self.gpu_ids[0]) 138 | else: 139 | torch.save(net.cpu().state_dict(), save_path) 140 | 141 | def load_networks(self, epoch): 142 | """ 143 | Load networks at specified epoch from the disk. 144 | 145 | Parameters: 146 | epoch (str) -- Which epoch to load 147 | """ 148 | for model_name in self.model_names: 149 | if isinstance(model_name, str): 150 | load_filename = '{:s}_net_{:s}.pth'.format(epoch, model_name) 151 | load_path = os.path.join(self.load_net_dir, load_filename) 152 | # Use the str to get the attribute aka the network (self.netG / self.netD) 153 | net = getattr(self, 'net' + model_name) 154 | # If we use multi GPUs and apply the data parallel 155 | if isinstance(net, torch.nn.DataParallel): 156 | net = net.module 157 | print('Loading the model from %s' % load_path) 158 | state_dict = torch.load(load_path, map_location=self.device) 159 | if hasattr(state_dict, '_metadata'): 160 | del state_dict._metadata 161 | 162 | net.load_state_dict(state_dict) 163 | 164 | def print_networks(self, param): 165 | """ 166 | Print the total number of parameters in the network and network architecture if detail is true 167 | Save the networks information to the disk 168 | """ 169 | message = '\n----------------------Networks Information----------------------' 170 | for model_name in self.model_names: 171 | if isinstance(model_name, str): 172 | net = getattr(self, 'net' + model_name) 173 | num_params = 0 174 | for parameter in net.parameters(): 175 | num_params += parameter.numel() 176 | if param.detail: 177 | message += '\n' + str(net) 178 | message += '\n[Network {:s}] Total number of parameters : {:.3f} M'.format(model_name, num_params / 1e6) 179 | message += '\n----------------------------------------------------------------\n' 180 | 181 | # Save the networks information to the disk 182 | net_info_filename = os.path.join(param.checkpoints_dir, param.experiment_name, 'net_info.txt') 183 | with open(net_info_filename, 'w') as log_file: 184 | log_file.write(message) 185 | 186 | print(message) 187 | 188 | def update_learning_rate(self): 189 | """ 190 | Update learning rates for all the networks 191 | Called at the end of each epoch 192 | """ 193 | 194 | lr = self.optimizers[0].param_groups[0]['lr'] 195 | # print('Learning rate for this epoch: %.7f' % lr) 196 | 197 | for scheduler in self.schedulers: 198 | if self.param.lr_policy == 'plateau': 199 | scheduler.step(self.metric) 200 | else: 201 | scheduler.step() 202 | 203 | def set_train(self): 204 | """ 205 | Set train mode for networks 206 | """ 207 | for model_name in self.model_names: 208 | if isinstance(model_name, str): 209 | # Use the str to get the attribute aka the network (self.netG / self.netD) 210 | net = getattr(self, 'net' + model_name) 211 | net.train() 212 | 213 | def set_eval(self): 214 | """ 215 | Set eval mode for networks 216 | """ 217 | for model_name in self.model_names: 218 | if isinstance(model_name, str): 219 | # Use the str to get the attribute aka the network (self.netG / self.netD) 220 | net = getattr(self, 'net' + model_name) 221 | net.eval() 222 | 223 | def test(self): 224 | """ 225 | Forward in testing to get the fake omics data 226 | """ 227 | with torch.no_grad(): 228 | self.forward() 229 | 230 | def init_losses_dict(self): 231 | """ 232 | initialize a losses dictionary 233 | """ 234 | losses_acc = OrderedDict() 235 | for name in self.loss_names: 236 | if isinstance(name, str): 237 | losses_acc[name] = [] 238 | return losses_acc 239 | 240 | def update_losses_dict(self, losses_dict): 241 | """ 242 | losses_dict (OrderedDict) -- the losses dictionary to be updated 243 | """ 244 | for name in self.loss_names: 245 | if isinstance(name, str): 246 | losses_dict[name].append(float(getattr(self, 'loss_' + name))) 247 | 248 | def init_metrics_dict(self): 249 | """ 250 | initialize a metrics dictionary 251 | """ 252 | metrics_acc = OrderedDict() 253 | for name in self.metric_names: 254 | if isinstance(name, str): 255 | metrics_acc[name] = [] 256 | return metrics_acc 257 | 258 | def update_metrics_dict(self, metrics_dict): 259 | """ 260 | Return performance metrics during testing 261 | metrics_dict (OrderedDict) -- the metrics dictionary to be updated 262 | """ 263 | for name in self.metric_names: 264 | if isinstance(name, str): 265 | metrics_dict[name].extend(getattr(self, 'metric_' + name)) 266 | 267 | def calculate_current_metrics(self): 268 | """ 269 | Calculate current metrics 270 | """ 271 | with torch.no_grad(): 272 | if self.gpu_ids: 273 | fake_A_array = self.fake_A_tensor.cpu().numpy() 274 | real_A_array = self.real_A_tensor.cpu().numpy() 275 | else: 276 | fake_A_array = self.fake_A_tensor.detach().numpy() 277 | real_A_array = self.real_A_tensor.detach().numpy() 278 | 279 | self.metric_MSE = [] 280 | self.metric_RMSE = [] 281 | self.metric_MAE = [] 282 | self.metric_MEDAE = [] 283 | self.metric_R2 = [] 284 | 285 | for i in range(fake_A_array.shape[0]): 286 | fake_A_single = fake_A_array[i].ravel() 287 | real_A_single = real_A_array[i].ravel() 288 | 289 | self.metric_MSE.append(metrics.mean_squared_error(real_A_single, fake_A_single)) 290 | self.metric_RMSE.append(metrics.mean_squared_error(real_A_single, fake_A_single, squared=False)) 291 | self.metric_MAE.append(metrics.mean_absolute_error(real_A_single, fake_A_single)) 292 | self.metric_MEDAE.append(metrics.median_absolute_error(real_A_single, fake_A_single)) 293 | self.metric_R2.append(metrics.r2_score(real_A_single, fake_A_single)) 294 | 295 | def init_fake_dict(self): 296 | """ 297 | initialize and return an empty fake array and an empty index array 298 | """ 299 | fake_dict = OrderedDict() 300 | fake_dict['index'] = np.zeros(shape=[0]) 301 | fake_dict['fake'] = np.zeros(shape=[0, self.param.A_dim]) 302 | return fake_dict 303 | 304 | def update_fake_dict(self, fake_dict): 305 | """ 306 | update the fake array that stores the predicted omics data 307 | fake_dict (OrderedDict) -- the fake array that stores the predicted omics data and the index array 308 | """ 309 | with torch.no_grad(): 310 | if self.param.add_channel: 311 | current_fake_array = np.squeeze(self.fake_A_tensor.cpu().numpy(), axis=1) 312 | else: 313 | current_fake_array = self.fake_A_tensor.cpu().numpy() 314 | fake_dict['fake'] = np.concatenate((fake_dict['fake'], current_fake_array)) 315 | current_index_array = self.data_index.cpu().numpy() 316 | fake_dict['index'] = np.concatenate((fake_dict['index'], current_index_array)) 317 | return fake_dict 318 | -------------------------------------------------------------------------------- /models/c_gan_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .basic_model import BasicModel 3 | from . import networks 4 | from . import losses 5 | import torch.nn as nn 6 | 7 | 8 | class CGanModel(BasicModel): 9 | """ 10 | This class implements the conditional GAN model, for learning a mapping from input omics data type to output omics 11 | data type given paired data. 12 | """ 13 | 14 | @staticmethod 15 | def modify_commandline_parameters(parser, is_train=True): 16 | # changing the default values of parameters to match the conditional GAN model 17 | parser.set_defaults(netG='fcg_sep', netD='fcd_sep') 18 | parser.add_argument('--latent_dim', type=int, default=256, 19 | help='the dimensionality of the latent space') 20 | # changing the default values of parameters to match the conditional GAN model 21 | if is_train: 22 | parser.add_argument('--lambda_dist', type=float, default=100.0, help='weight for the dist loss') 23 | return parser 24 | 25 | def __init__(self, param): 26 | """ 27 | Initialize the conditional GAN class. 28 | """ 29 | BasicModel.__init__(self, param) 30 | # specify the training losses you want to print out. 31 | self.loss_names = ['D_GAN', 'G_GAN', 'G_dist'] 32 | # specify the models you want to save to the disk. 33 | if self.isTrain: 34 | self.model_names = ['G', 'D'] 35 | else: # during test phase, only load G 36 | self.model_names = ['G'] 37 | 38 | # define networks (both generator and discriminator) 39 | self.netG = networks.define_G(param.input_chan_num, param.output_chan_num, param.netG, param.A_dim, param.B_dim, 40 | param.gen_filter_num, param.conv_k_size, param.norm_type, param.init_type, 41 | param.init_gain, self.gpu_ids, param.leaky_slope, param.dropout_p, param.latent_dim) 42 | 43 | # define a discriminator if it is the training phase, if it's the testing phase the discriminator is not necessary 44 | if self.isTrain: 45 | self.netD = networks.define_D(param.input_chan_num, param.output_chan_num, param.dis_filter_num, param.netD, 46 | param.A_dim, param.B_dim, param.layer_num_D, param.norm_type, param.init_type, 47 | param.init_gain, self.gpu_ids, param.leaky_slope, param.dropout_p) 48 | 49 | # define loss functions: G = L_GAN + λ L_dist 50 | # The GAN part of the loss function, this return a loss function nn.Module not a value 51 | # self.device was defined in BaseModel 52 | self.lossFuncGAN = losses.GANLossObj(param.GAN_mode).to(self.device) 53 | self.lossFuncDist = losses.get_dist_loss(param.dist_loss) 54 | 55 | # Set optimizer for both generator and discriminator 56 | # generator and discriminator actually can set to different initial learning rate 57 | self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=param.lr_G, betas=(param.beta1, 0.999)) 58 | self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=param.lr_D, betas=(param.beta1, 0.999)) 59 | # optimizer list was already defined in BaseModel 60 | self.optimizers.append(self.optimizer_G) 61 | self.optimizers.append(self.optimizer_D) 62 | 63 | self.loss_D_fake = [] 64 | self.loss_D_real = [] 65 | self.loss_D_GAN = [] 66 | self.loss_G_GAN = [] 67 | self.loss_G_dist = [] 68 | self.loss_G = [] 69 | 70 | if self.param.zo_norm: 71 | self.sigmoid = nn.Sigmoid() 72 | 73 | def set_input(self, input_dict): 74 | """ 75 | Unpack input data from the output dictionary of the dataloader 76 | 77 | Parameters: 78 | input_dict (dict): include the data tensor and its index. 79 | """ 80 | self.real_A_tensor = input_dict['A_tensor'].to(self.device) 81 | if self.param.ch_separate: 82 | self.real_B_tensor = [] 83 | for ch in range(0, 23): 84 | self.real_B_tensor.append(input_dict['B_tensor'][ch].to(self.device)) 85 | else: 86 | self.real_B_tensor = input_dict['B_tensor'].to(self.device) 87 | self.data_index = input_dict['index'] 88 | 89 | def forward(self): 90 | # Default B -> A 91 | self.fake_A_tensor = self.netG(self.real_B_tensor) # A' = G(B) 92 | if self.param.zo_norm: 93 | self.fake_A_tensor = self.sigmoid(self.fake_A_tensor) 94 | self.fake_A_tensor = (self.param.target_max - self.param.target_min) * self.fake_A_tensor + self.param.target_min 95 | # Calculate metrics of the fake tensor 96 | self.calculate_current_metrics() 97 | 98 | def backward_G(self): 99 | """Calculate GAN and dist loss for the generator""" 100 | # G(B) should fake the discriminator to treat it as real omics data 101 | # The different part compared with the backward_D is that we don't need the detach fake tensor here 102 | pred_fake = self.netD(self.fake_A_tensor, self.real_B_tensor) # The prediction vector get from the discriminator for the fake omics data 103 | self.loss_G_GAN = self.lossFuncGAN(pred_fake, True) # The boolean variable will be extend to a vector as the same size of pred_fake 104 | 105 | # G(B) should be as close as A, we use the distance loss 106 | if self.param.zo_norm and self.param.dist_loss == 'BCE': 107 | self.loss_G_dist = self.lossFuncDist((self.fake_A_tensor - self.param.target_min)/(self.param.target_max - self.param.target_min), (self.real_A_tensor - self.param.target_min)/(self.param.target_max - self.param.target_min)) 108 | else: 109 | self.loss_G_dist = self.lossFuncDist(self.fake_A_tensor, self.real_A_tensor) 110 | 111 | # Combine the loss and calculate gradients 112 | # G = L_GAN + λ L_dist 113 | # The parameter lambda_dist was introduced in this class 114 | self.loss_G = self.loss_G_GAN + self.param.lambda_dist * self.loss_G_dist 115 | self.loss_G.backward() 116 | 117 | def backward_D(self): 118 | """Calculate GAN loss for the discriminator""" 119 | # Fake 120 | # Stop backprop to the generator by detaching fake_A 121 | # Conditional GAN was applied so both the input and output of generator were fed to discriminator 122 | pred_fake = self.netD(self.fake_A_tensor.detach(), self.real_B_tensor) # the prediction vector get from the discriminator for the fake omics data 123 | self.loss_D_fake = self.lossFuncGAN(pred_fake, False) # The boolean variable will be extend to a vector as the same size of pred_fake 124 | 125 | # Real 126 | pred_real = self.netD(self.real_A_tensor, self.real_B_tensor) 127 | self.loss_D_real = self.lossFuncGAN(pred_real, True) 128 | 129 | # Combine the loss and calculate gradients 130 | self.loss_D_GAN = (self.loss_D_fake + self.loss_D_real) / 2 131 | self.loss_D_GAN.backward() 132 | 133 | def update(self): 134 | self.forward() # Get the fake omics data: G(B) 135 | 136 | # Update parameters of the discriminator 137 | # the method is defined in BaseModel 138 | self.set_requires_grad(self.netD, True) # Enable backprop for D 139 | self.optimizer_D.zero_grad() # Set D's gradients to zero 140 | self.backward_D() # Calculate gradients for D 141 | self.optimizer_D.step() # Update D's weights 142 | 143 | # Update parameters of the generator 144 | self.set_requires_grad(self.netD, False) # Stop backprop for D when optimizing G 145 | self.optimizer_G.zero_grad() # Set G's gradients to zero 146 | self.backward_G() # Calculate gradients for G 147 | self.optimizer_G.step() # Update G's weights 148 | 149 | -------------------------------------------------------------------------------- /models/c_gan_ved_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .basic_model import BasicModel 3 | from . import networks 4 | from . import losses 5 | import torch.nn as nn 6 | 7 | 8 | class CGanVedModel(BasicModel): 9 | """ 10 | This class implements the conditional GAN VED model, for learning a mapping from input omics data type to output omics 11 | data type given paired data. 12 | """ 13 | @staticmethod 14 | def modify_commandline_parameters(parser, is_train=True): 15 | # changing the default values of parameters to match the conditional GAN model 16 | parser.set_defaults(netG='fc_ved_sep', netD='fcd_sep') 17 | parser.add_argument('--latent_dim', type=int, default=256, 18 | help='the dimensionality of the latent space') 19 | if is_train: 20 | parser.add_argument('--lambda_dist', type=float, default=100.0, 21 | help='weight for the dist loss') 22 | parser.add_argument('--lambda_kl', type=float, default=1.0, 23 | help='weight for the kl loss') 24 | return parser 25 | 26 | def __init__(self, param): 27 | """ 28 | Initialize the conditional GAN VED class. 29 | """ 30 | BasicModel.__init__(self, param) 31 | # Declare mean and var for the kl 32 | self.mean = [] 33 | self.log_var = [] 34 | # specify the training losses you want to print out. 35 | self.loss_names = ['D_GAN', 'G_GAN', 'G_dist', 'G_kl'] 36 | # specify the models you want to save to the disk. 37 | if self.isTrain: 38 | self.model_names = ['G', 'D'] 39 | else: # during test phase, only load G 40 | self.model_names = ['G'] 41 | 42 | # define networks (both generator and discriminator) 43 | self.netG = networks.define_G(param.input_chan_num, param.output_chan_num, param.netG, param.A_dim, param.B_dim, 44 | param.gen_filter_num, param.conv_k_size, param.norm_type, param.init_type, 45 | param.init_gain, self.gpu_ids, param.leaky_slope, param.dropout_p, param.latent_dim) 46 | 47 | # define a discriminator if it is the training phase, if it's the testing phase the discriminator is not necessary 48 | if self.isTrain: 49 | self.netD = networks.define_D(param.input_chan_num, param.output_chan_num, param.dis_filter_num, param.netD, 50 | param.A_dim, param.B_dim, param.layer_num_D, param.norm_type, param.init_type, 51 | param.init_gain, self.gpu_ids, param.leaky_slope, param.dropout_p) 52 | 53 | # define loss functions: G = L_GAN + λ1 L_dist + λ2 KL 54 | # The GAN part of the loss function, this return a loss function nn.Module not a value 55 | # self.device was defined in BaseModel 56 | self.lossFuncGAN = losses.GANLossObj(param.GAN_mode).to(self.device) 57 | self.lossFuncDist = losses.get_dist_loss(param.dist_loss) 58 | 59 | # Set optimizer for both generator and discriminator 60 | # generator and discriminator actually can set to different initial learning rate 61 | self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=param.lr_G, betas=(param.beta1, 0.999)) 62 | self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=param.lr_D, betas=(param.beta1, 0.999)) 63 | # optimizer list was already defined in BaseModel 64 | self.optimizers.append(self.optimizer_G) 65 | self.optimizers.append(self.optimizer_D) 66 | 67 | self.loss_D_fake = [] 68 | self.loss_D_real = [] 69 | self.loss_D_GAN = [] 70 | self.loss_G_GAN = [] 71 | self.loss_G_dist = [] 72 | self.loss_G_kl = [] 73 | self.loss_G = [] 74 | 75 | if self.param.zo_norm: 76 | self.sigmoid = nn.Sigmoid() 77 | 78 | def set_input(self, input_dict): 79 | """ 80 | Unpack input data from the output dictionary of the dataloader 81 | 82 | Parameters: 83 | input_dict (dict): include the data tensor and its index. 84 | """ 85 | self.real_A_tensor = input_dict['A_tensor'].to(self.device) 86 | if self.param.ch_separate: 87 | self.real_B_tensor = [] 88 | for ch in range(0, 23): 89 | self.real_B_tensor.append(input_dict['B_tensor'][ch].to(self.device)) 90 | else: 91 | self.real_B_tensor = input_dict['B_tensor'].to(self.device) 92 | self.data_index = input_dict['index'] 93 | 94 | def forward(self): 95 | # Default B -> A 96 | _, self.fake_A_tensor, self.mean, self.log_var = self.netG(self.real_B_tensor) # A' = G(B) 97 | if self.param.zo_norm: 98 | self.fake_A_tensor = self.sigmoid(self.fake_A_tensor) 99 | self.fake_A_tensor = (self.param.target_max - self.param.target_min) * self.fake_A_tensor + self.param.target_min 100 | # Calculate metrics of the fake tensor 101 | self.calculate_current_metrics() 102 | 103 | def backward_G(self): 104 | """Calculate GAN and dist loss for the generator""" 105 | # G(B) should fake the discriminator to treat it as real omics data 106 | # The different part compared with the backward_D is that we don't need the detach fake tensor here 107 | pred_fake = self.netD(self.fake_A_tensor, self.real_B_tensor) # The prediction vector get from the discriminator for the fake omics data 108 | self.loss_G_GAN = self.lossFuncGAN(pred_fake, True) # The boolean variable will be extend to a vector as the same size of pred_fake 109 | 110 | # G(B) should be as close as A, we use the distance loss 111 | if self.param.zo_norm and self.param.dist_loss == 'BCE': 112 | self.loss_G_dist = self.lossFuncDist( 113 | (self.fake_A_tensor - self.param.target_min) / (self.param.target_max - self.param.target_min), 114 | (self.real_A_tensor - self.param.target_min) / (self.param.target_max - self.param.target_min)) 115 | else: 116 | self.loss_G_dist = self.lossFuncDist(self.fake_A_tensor, self.real_A_tensor) 117 | 118 | # Add the kl constrain 119 | self.loss_G_kl = losses.kl_loss(self.mean, self.log_var) 120 | 121 | # Combine the loss and calculate gradients 122 | # G = L_GAN + λ1 L_dist + λ2 KL 123 | # The parameter lambda_dist was introduced in this class 124 | self.loss_G = self.loss_G_GAN + self.param.lambda_dist * self.loss_G_dist + self.param.lambda_kl * self.loss_G_kl 125 | self.loss_G.backward() 126 | 127 | def backward_D(self): 128 | """Calculate GAN loss for the discriminator""" 129 | # Fake 130 | # Stop backprop to the generator by detaching fake_A 131 | # Conditional GAN was applied so both the input and output of generator were fed to discriminator 132 | pred_fake = self.netD(self.fake_A_tensor.detach(), self.real_B_tensor) # the prediction vector get from the discriminator for the fake omics data 133 | self.loss_D_fake = self.lossFuncGAN(pred_fake, False) # The boolean variable will be extend to a vector as the same size of pred_fake 134 | 135 | # Real 136 | pred_real = self.netD(self.real_A_tensor, self.real_B_tensor) 137 | self.loss_D_real = self.lossFuncGAN(pred_real, True) 138 | 139 | # Combine the loss and calculate gradients 140 | self.loss_D_GAN = (self.loss_D_fake + self.loss_D_real) / 2 141 | self.loss_D_GAN.backward() 142 | 143 | def update(self): 144 | self.forward() # Get the fake omics data: G(B) 145 | 146 | # Update parameters of the discriminator 147 | # the method is defined in BaseModel 148 | self.set_requires_grad(self.netD, True) # Enable backprop for D 149 | self.optimizer_D.zero_grad() # Set D's gradients to zero 150 | self.backward_D() # Calculate gradients for D 151 | self.optimizer_D.step() # Update D's weights 152 | 153 | # Update parameters of the generator 154 | self.set_requires_grad(self.netD, False) # Stop backprop for D when optimizing G 155 | self.optimizer_G.zero_grad() # Set G's gradients to zero 156 | self.backward_G() # Calculate gradients for G 157 | self.optimizer_G.step() # Update G's weights 158 | 159 | -------------------------------------------------------------------------------- /models/losses.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class GANLossObj(nn.Module): 6 | """ 7 | Define different GAN objectives. 8 | 9 | The GANLoss class abstracts away the need to create the target label tensor that has the same size as the input. 10 | Because we need to calculate the loss between the output vector of the discriminator and True/False, this class 11 | automatically extend True/False to an all_1 vector or an all_0 vector so the loss function is able to calculate. 12 | """ 13 | def __init__(self, GAN_mode, real_label=1.0, fake_label=0.0): 14 | """ 15 | Initialize the GANLossObj class. 16 | 17 | Parameters: 18 | GAN_mode (str) -- The type of GAN objective. Vanilla GAN loss is the cross-entropy objective used in the original GAN paper. 19 | real_label (float) -- the label value for the real omics data 20 | fake_label (float) -- the label value for the synthetic omics data 21 | 22 | Note: 23 | Do not use sigmoid as the last layer of Discriminator. 24 | LSGAN needs no sigmoid. Vanilla GANs will handle it with BCEWithLogitsLoss. 25 | This is why we don't put a sigmoid layer in the discriminator 26 | """ 27 | super(GANLossObj, self).__init__() 28 | # registering the labels for model storage 29 | self.register_buffer('real_label', torch.tensor(real_label)) 30 | self.register_buffer('fake_label', torch.tensor(fake_label)) 31 | self.GAN_mode = GAN_mode 32 | # Determine the loss function 33 | if GAN_mode == 'vanilla': 34 | self.loss = nn.BCEWithLogitsLoss() 35 | elif GAN_mode == 'lsgan': 36 | self.loss = nn.MSELoss() 37 | elif GAN_mode == 'wgangp': 38 | self.loss = None 39 | else: 40 | raise NotImplementedError('GAN mode %s is not found' % GAN_mode) 41 | 42 | def get_target_tensor(self, prediction, target_is_real): 43 | """ 44 | Create label tensors with the same size as the input tensor 45 | 46 | Parameters: 47 | prediction (tensor) -- the output from a discriminator 48 | target_is_real (bool) -- whether the ground truth label is for real or fake 49 | 50 | Returns: 51 | A label tensor filled with ground truth label with the size of the input tensor 52 | """ 53 | 54 | if target_is_real: 55 | target_tensor = self.real_label 56 | else: 57 | target_tensor = self.fake_label 58 | return target_tensor.expand_as(prediction) 59 | 60 | def __call__(self, prediction, target_is_real): 61 | """ 62 | Calculate the loss given output from the discriminator and the grount truth tensor. 63 | 64 | Parameters: 65 | prediction (tensor) -- the output from a discriminator 66 | target_is_real (bool) -- whether the ground truth label is for real or fake 67 | 68 | Returns: 69 | The calculated loss 70 | """ 71 | if self.GAN_mode in ['lsgan', 'vanilla']: 72 | target_tensor = self.get_target_tensor(prediction, target_is_real) 73 | loss = self.loss(prediction, target_tensor) 74 | elif self.GAN_mode == 'wgangp': 75 | if target_is_real: 76 | loss = -prediction.mean() 77 | else: 78 | loss = prediction.mean() 79 | else: 80 | raise NotImplementedError('GAN mode %s is not found' % self.GAN_mode) 81 | return loss 82 | 83 | 84 | def get_dist_loss(loss_name, reduction='mean'): 85 | """ 86 | Return the distance loss function. 87 | Parameters: 88 | loss_name (str) -- the name of the loss function: BCE | MSE | L1 | CE 89 | reduction (str) -- the reduction method applied to the loss function: sum | mean 90 | """ 91 | if loss_name == 'BCE': 92 | return nn.BCELoss(reduction=reduction) 93 | elif loss_name == 'MSE': 94 | return nn.MSELoss(reduction=reduction) 95 | elif loss_name == 'L1': 96 | return nn.L1Loss(reduction=reduction) 97 | elif loss_name == 'CE': 98 | return nn.CrossEntropyLoss(reduction=reduction) 99 | else: 100 | raise NotImplementedError('Loss function %s is not found' % loss_name) 101 | 102 | 103 | def kl_loss(mean, log_var, reduction='mean'): 104 | part_loss = 1 + log_var - mean.pow(2) - log_var.exp() 105 | if reduction == 'mean': 106 | loss = -0.5 * torch.mean(part_loss) 107 | else: 108 | loss = -0.5 * torch.sum(part_loss) 109 | return loss -------------------------------------------------------------------------------- /models/networks.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import functools 4 | from torch.nn import init 5 | from torch.optim import lr_scheduler 6 | 7 | 8 | class Identity(nn.Module): 9 | def forward(self, x): 10 | return x 11 | 12 | 13 | # GENERATOR 14 | 15 | # For one dimensional U-Net generator 16 | class UNetDe(nn.Module): 17 | """ 18 | Create a 1D U-Net network (decreasing skip connection) 19 | """ 20 | def __init__(self, input_chan_num, output_chan_num, output_dim, filter_num=64, kernel_size=9, 21 | norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0): 22 | """ 23 | Construct a one dimensional U-Net 24 | Parameters: 25 | input_chan_num (int) -- the number of channels in input omics data 26 | output_chan_num (int) -- the number of channels in output omics data 27 | output_dim (int) -- the dimension of the output omics data 28 | filter_num (int) -- the number of filters in the first convolution layer 29 | kernel_size (int) -- the kernel size of convolution layers 30 | norm_layer -- normalization layer 31 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 32 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 33 | """ 34 | super(UNetDe, self).__init__() 35 | self.input_chan_num = input_chan_num 36 | self.output_chan_num = output_chan_num 37 | self.output_dim = output_dim 38 | 39 | # Encoder 40 | # 1 -> 64 deal with the input data the first double convolution layer 41 | self.input_conv = DoubleConv1D(input_chan_num, filter_num, kernel_size=kernel_size, norm_layer=norm_layer, 42 | leaky_slope=leaky_slope) 43 | # 64 -> 128 44 | self.down_sample_1 = DownSample(filter_num, filter_num * 2, down_ratio=4, kernel_size=kernel_size, 45 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=0) 46 | # 128 -> 256 47 | self.down_sample_2 = DownSample(filter_num * 2, filter_num * 4, down_ratio=4, kernel_size=kernel_size, 48 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 49 | # 256 -> 512 50 | self.down_sample_3 = DownSample(filter_num * 4, filter_num * 8, down_ratio=4, kernel_size=kernel_size, 51 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=0) 52 | # 512 -> 1024 53 | self.down_sample_4 = DownSample(filter_num * 8, filter_num * 16, down_ratio=4, kernel_size=kernel_size, 54 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 55 | 56 | # Decoder 57 | # 1024 -> 512 58 | self.up_sample_1 = UpSampleDe(filter_num * 16, filter_num * 8, up_ratio=2, skip_ratio=2, kernel_size=kernel_size, 59 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=0) 60 | # 512 -> 256 61 | self.up_sample_2 = UpSampleDe(filter_num * 8, filter_num * 4, up_ratio=2, skip_ratio=4, kernel_size=kernel_size, 62 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 63 | # 256 -> 128 64 | self.up_sample_3 = UpSampleDe(filter_num * 4, filter_num * 2, up_ratio=2, skip_ratio=8, kernel_size=kernel_size, 65 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=0) 66 | # 128 -> 64 67 | self.up_sample_4 = UpSampleDe(filter_num * 2, filter_num, up_ratio=5, skip_ratio=6, kernel_size=kernel_size, 68 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 69 | # 64 -> 1 70 | self.output_conv = OutputConv(filter_num, output_chan_num) 71 | 72 | def forward(self, x): 73 | x1 = self.input_conv(x) 74 | x2 = self.down_sample_1(x1) 75 | x3 = self.down_sample_2(x2) 76 | x4 = self.down_sample_3(x3) 77 | x5 = self.down_sample_4(x4) 78 | y1 = self.up_sample_1(x5, x4) 79 | y2 = self.up_sample_2(y1, x3) 80 | y3 = self.up_sample_3(y2, x2) 81 | y4 = self.up_sample_4(y3, x1) 82 | y5 = self.output_conv(y4) 83 | # Let the output dim the same as targeting dim 84 | output = y5[:, :, 0:self.output_dim] 85 | return output 86 | 87 | 88 | class UNetIn(nn.Module): 89 | """ 90 | Create a 1D U-Net network 91 | """ 92 | def __init__(self, input_chan_num, output_chan_num, output_dim, filter_num=64, kernel_size=9, 93 | norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0): 94 | """ 95 | Construct a one dimensional U-Net 96 | Parameters: 97 | input_chan_num (int) -- the number of channels in input omics data 98 | output_chan_num (int) -- the number of channels in output omics data 99 | output_dim (int) -- the dimension of the output omics data 100 | filter_num (int) -- the number of filters in the first convolution layer 101 | kernel_size (int) -- the kernel size of convolution layers 102 | norm_layer -- normalization layer 103 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 104 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 105 | """ 106 | super(UNetIn, self).__init__() 107 | self.input_chan_num = input_chan_num 108 | self.output_chan_num = output_chan_num 109 | self.output_dim = output_dim 110 | 111 | # Encoder 112 | # 1 -> 64 deal with the input data the first double convolution layer 113 | self.input_conv = DoubleConv1D(input_chan_num, filter_num, kernel_size=kernel_size, norm_layer=norm_layer, 114 | leaky_slope=leaky_slope) 115 | # 64 -> 128 116 | self.down_sample_1 = DownSample(filter_num, filter_num * 2, down_ratio=2, kernel_size=kernel_size, 117 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=0) 118 | # 128 -> 256 119 | self.down_sample_2 = DownSample(filter_num * 2, filter_num * 4, down_ratio=2, kernel_size=kernel_size, 120 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 121 | # 256 -> 512 122 | self.down_sample_3 = DownSample(filter_num * 4, filter_num * 8, down_ratio=2, kernel_size=kernel_size, 123 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=0) 124 | # 512 -> 1024 125 | self.down_sample_4 = DownSample(filter_num * 8, filter_num * 16, down_ratio=2, kernel_size=kernel_size, 126 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 127 | 128 | # Decoder 129 | # 1024 -> 512 130 | self.up_sample_1 = UpSampleIn(filter_num * 16, filter_num * 8, up_ratio=4, skip_ratio=2, kernel_size=kernel_size, 131 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=0) 132 | # 512 -> 256 133 | self.up_sample_2 = UpSampleIn(filter_num * 8, filter_num * 4, up_ratio=2, skip_ratio=2, kernel_size=kernel_size, 134 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 135 | # 256 -> 128 136 | self.up_sample_3 = UpSampleIn(filter_num * 4, filter_num * 2, up_ratio=2, skip_ratio=2, kernel_size=kernel_size, 137 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=0) 138 | # 128 -> 64 139 | self.up_sample_4 = UpSampleIn(filter_num * 2, filter_num, up_ratio=2, skip_ratio=2, kernel_size=kernel_size, 140 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 141 | # 64 -> 1 142 | self.output_conv = OutputConv(filter_num, output_chan_num) 143 | 144 | def forward(self, x): 145 | x1 = self.input_conv(x) 146 | x2 = self.down_sample_1(x1) 147 | x3 = self.down_sample_2(x2) 148 | x4 = self.down_sample_3(x3) 149 | x5 = self.down_sample_4(x4) 150 | y1 = self.up_sample_1(x5, x4) 151 | y2 = self.up_sample_2(y1, x3) 152 | y3 = self.up_sample_3(y2, x2) 153 | y4 = self.up_sample_4(y3, x1) 154 | y5 = self.output_conv(y4) 155 | # Let the output dim the same as targeting dim 156 | output = y5[:, :, 0:self.output_dim] 157 | return output 158 | 159 | 160 | class UNetSDe(nn.Module): 161 | """ 162 | Create a simplified 1D U-Net network (decreasing skip connection) 163 | """ 164 | def __init__(self, input_chan_num, output_chan_num, output_dim, filter_num=64, kernel_size=9, 165 | norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0): 166 | """ 167 | Construct a one dimensional U-Net 168 | Parameters: 169 | input_chan_num (int) -- the number of channels in input omics data 170 | output_chan_num (int) -- the number of channels in output omics data 171 | output_dim (int) -- the dimension of the output omics data 172 | filter_num (int) -- the number of filters in the first convolution layer 173 | kernel_size (int) -- the kernel size of convolution layers 174 | norm_layer -- normalization layer 175 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 176 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 177 | """ 178 | super(UNetSDe, self).__init__() 179 | self.input_chan_num = input_chan_num 180 | self.output_chan_num = output_chan_num 181 | self.output_dim = output_dim 182 | 183 | # Encoder 184 | # 1 -> 64 deal with the input data the first double convolution layer 185 | self.input_conv = DoubleConv1D(input_chan_num, filter_num, kernel_size=kernel_size, norm_layer=norm_layer, 186 | leaky_slope=leaky_slope) 187 | # 64 -> 256 188 | self.down_sample_1 = DownSample(filter_num, filter_num * 4, down_ratio=16, kernel_size=kernel_size, 189 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 190 | # 256 -> 1024 191 | self.down_sample_2 = DownSample(filter_num * 4, filter_num * 16, down_ratio=16, kernel_size=kernel_size, 192 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 193 | 194 | # Decoder 195 | # 1024 -> 256 196 | self.up_sample_1 = UpSampleDe(filter_num * 16, filter_num * 4, up_ratio=4, skip_ratio=4, kernel_size=kernel_size, 197 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 198 | # 256 -> 64 199 | self.up_sample_2 = UpSampleDe(filter_num * 4, filter_num, up_ratio=10, skip_ratio=6, kernel_size=kernel_size, 200 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 201 | # 64 -> 1 202 | self.output_conv = OutputConv(filter_num, output_chan_num) 203 | 204 | def forward(self, x): 205 | x1 = self.input_conv(x) 206 | x2 = self.down_sample_1(x1) 207 | x3 = self.down_sample_2(x2) 208 | y1 = self.up_sample_1(x3, x2) 209 | y2 = self.up_sample_2(y1, x1) 210 | y3 = self.output_conv(y2) 211 | # Let the output dim the same as targeting dim 212 | output = y3[:, :, 0:self.output_dim] 213 | return output 214 | 215 | 216 | class DoubleConv1D(nn.Module): 217 | """ 218 | (Convolution1D => Norm1D => LeakyReLU) * 2 219 | The omics data dimension keep the same during this process 220 | """ 221 | def __init__(self, input_chan_num, output_chan_num, kernel_size=9, norm_layer=nn.BatchNorm1d, leaky_slope=0.2): 222 | """ 223 | Construct a double convolution block 224 | Parameters: 225 | input_chan_num (int) -- the number of channels of the input tensor 226 | output_chan_num (int) -- the number of channels of the output tensor 227 | kernel_size (int) -- the kernel size of the convolution layer 228 | norm_layer -- normalization layer 229 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 230 | """ 231 | super(DoubleConv1D, self).__init__() 232 | 233 | # Only if the norm method is instance norm we use bias for the corresponding conv layer 234 | if type(norm_layer) == functools.partial: 235 | use_bias = norm_layer.func == nn.InstanceNorm1d 236 | else: 237 | use_bias = norm_layer == nn.InstanceNorm1d 238 | 239 | self.double_conv_1d = nn.Sequential( 240 | nn.Conv1d(input_chan_num, output_chan_num, kernel_size=kernel_size, padding=kernel_size // 2, bias=use_bias), 241 | norm_layer(output_chan_num), 242 | nn.LeakyReLU(negative_slope=leaky_slope, inplace=True), 243 | nn.Conv1d(output_chan_num, output_chan_num, kernel_size=kernel_size, padding=kernel_size // 2, bias=use_bias), 244 | norm_layer(output_chan_num), 245 | nn.LeakyReLU(negative_slope=leaky_slope, inplace=True) 246 | ) 247 | 248 | def forward(self, x): 249 | return self.double_conv_1d(x) 250 | 251 | 252 | class DownSample(nn.Module): 253 | """ 254 | Downsampling with MaxPool then call the DoubleConv1D module 255 | The output dimension = input dimension // ratio 256 | """ 257 | def __init__(self, input_chan_num, output_chan_num, down_ratio, kernel_size=9, norm_layer=nn.BatchNorm1d, 258 | leaky_slope=0.2, dropout_p=0): 259 | """ 260 | Construct a downsampling block 261 | Parameters: 262 | input_chan_num (int) -- the number of channels of the input tensor 263 | output_chan_num (int) -- the number of channels of the output tensor 264 | down_ratio (int) -- the kernel size and stride of the MaxPool1d layer 265 | kernel_size (int) -- the kernel size of the DoubleConv1D block 266 | norm_layer -- normalization layer 267 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 268 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 269 | """ 270 | super(DownSample, self).__init__() 271 | self.maxpool_double_conv = nn.Sequential( 272 | nn.MaxPool1d(down_ratio), 273 | nn.Dropout(p=dropout_p), 274 | DoubleConv1D(input_chan_num, output_chan_num, kernel_size, norm_layer, leaky_slope) 275 | ) 276 | 277 | def forward(self, x): 278 | return self.maxpool_double_conv(x) 279 | 280 | 281 | class UpSampleDe(nn.Module): 282 | """ 283 | Upsampling with ConvTranspose1d then call the DoubleConv1D module (decreasing skip connection) 284 | The output dimension = input dimension * ratio 285 | """ 286 | def __init__(self, input_chan_num, output_chan_num, up_ratio, skip_ratio, kernel_size=9, norm_layer=nn.BatchNorm1d, 287 | leaky_slope=0.2, dropout_p=0): 288 | """ 289 | Construct a upsampling block 290 | Parameters: 291 | input_chan_num (int) -- the number of channels of the input tensor (the tensor from get from the last layer, not the tensor from the skip-connection mechanism) 292 | output_chan_num (int) -- the number of channels of the output tensor 293 | up_ratio (int) -- the kernel size and stride of the ConvTranspose1d layer 294 | skip_ratio (int) -- the kernel size the MaxPool1d layer for the skip connection 295 | kernel_size (int) -- the kernel size of the DoubleConv1D block 296 | norm_layer -- normalization layer 297 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 298 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 299 | """ 300 | super(UpSampleDe, self).__init__() 301 | self.up_sample = nn.ConvTranspose1d(input_chan_num, output_chan_num, kernel_size=up_ratio, stride=up_ratio) 302 | self.skip_connect_pool = nn.MaxPool1d(skip_ratio) 303 | self.dropout = nn.Dropout(p=dropout_p) 304 | self.double_conv = DoubleConv1D(output_chan_num * 2, output_chan_num, kernel_size, norm_layer, leaky_slope) 305 | 306 | def forward(self, input1, input2): 307 | x1 = self.up_sample(input1) 308 | x2 = self.skip_connect_pool(input2) 309 | x2_crop = x2[:, :, 0:x1.shape[2]] 310 | # The skip connection mechanism 311 | x = torch.cat([x2_crop, x1], dim=1) 312 | x_dropout = self.dropout(x) 313 | return self.double_conv(x_dropout) 314 | 315 | 316 | class UpSampleIn(nn.Module): 317 | """ 318 | Upsampling with ConvTranspose1d then call the DoubleConv1D module (increasing skip connection) 319 | The output dimension = input dimension * ratio 320 | """ 321 | def __init__(self, input_chan_num, output_chan_num, up_ratio, skip_ratio, kernel_size=9, norm_layer=nn.BatchNorm1d, 322 | leaky_slope=0.2, dropout_p=0): 323 | """ 324 | Construct a upsampling block 325 | Parameters: 326 | input_chan_num (int) -- the number of channels of the input tensor (the tensor from get from the last layer, not the tensor from the skip-connection mechanism) 327 | output_chan_num (int) -- the number of channels of the output tensor 328 | up_ratio (int) -- the kernel size and stride of the ConvTranspose1d layer 329 | skip_ratio (int) -- the kernel size and stride of the ConvTranspose1d layer for the skip connection 330 | kernel_size (int) -- the kernel size of the DoubleConv1D block 331 | norm_layer -- normalization layer 332 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 333 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 334 | """ 335 | super(UpSampleIn, self).__init__() 336 | self.up_sample = nn.ConvTranspose1d(input_chan_num, output_chan_num, kernel_size=up_ratio, stride=up_ratio) 337 | self.skip_connect_pool = nn.ConvTranspose1d(output_chan_num, output_chan_num, kernel_size=skip_ratio, stride=skip_ratio) 338 | self.dropout = nn.Dropout(p=dropout_p) 339 | self.double_conv = DoubleConv1D(output_chan_num * 2, output_chan_num, kernel_size, norm_layer, leaky_slope) 340 | 341 | def forward(self, input1, input2): 342 | x1 = self.up_sample(input1) 343 | x2 = self.skip_connect_pool(input2) 344 | x2_crop = x2[:, :, 0:x1.shape[2]] 345 | # The skip connection mechanism 346 | x = torch.cat([x2_crop, x1], dim=1) 347 | x_dropout = self.dropout(x) 348 | return self.double_conv(x_dropout) 349 | 350 | 351 | class OutputConv(nn.Module): 352 | """ 353 | Output convolution layer 354 | """ 355 | def __init__(self, input_chan_num, output_chan_num): 356 | """ 357 | Construct the output convolution layer 358 | Parameters: 359 | input_chan_num (int) -- the number of channels of the input tensor 360 | output_chan_num (int) -- the number of channels of the output omics data 361 | """ 362 | super(OutputConv, self).__init__() 363 | self.output_conv = nn.Conv1d(input_chan_num, output_chan_num, kernel_size=1) 364 | 365 | def forward(self, x): 366 | return self.output_conv(x) 367 | 368 | 369 | class SingleConv1D(nn.Module): 370 | """ 371 | Convolution1D => Norm1D => LeakyReLU 372 | The omics data dimension keep the same during this process 373 | """ 374 | def __init__(self, input_chan_num, output_chan_num, kernel_size=9, norm_layer=nn.BatchNorm1d, leaky_slope=0.2): 375 | """ 376 | Construct a double convolution block 377 | Parameters: 378 | input_chan_num (int) -- the number of channels of the input tensor 379 | output_chan_num (int) -- the number of channels of the output tensor 380 | kernel_size (int) -- the kernel size of the convolution layer 381 | norm_layer -- normalization layer 382 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 383 | """ 384 | super(SingleConv1D, self).__init__() 385 | 386 | # Only if the norm method is instance norm we use bias for the corresponding conv layer 387 | if type(norm_layer) == functools.partial: 388 | use_bias = norm_layer.func == nn.InstanceNorm1d 389 | else: 390 | use_bias = norm_layer == nn.InstanceNorm1d 391 | 392 | self.double_conv_1d = nn.Sequential( 393 | nn.Conv1d(input_chan_num, output_chan_num, kernel_size=kernel_size, padding=kernel_size // 2, bias=use_bias), 394 | norm_layer(output_chan_num), 395 | nn.LeakyReLU(negative_slope=leaky_slope, inplace=True), 396 | ) 397 | 398 | def forward(self, x): 399 | return self.double_conv_1d(x) 400 | 401 | # For fully-connected generator 402 | class FcVed(nn.Module): 403 | """ 404 | Defines a fully-connected variational encoder-decoder model for DNA methylation to gene expression translation 405 | DNA methylation input not separated by chromosome 406 | """ 407 | def __init__(self, input_dim, output_dim, norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0, dim_1B=4096, dim_2B=1024, 408 | dim_3_B=512, dim_1A=4096, dim_2A=1024, dim_3_A=512, latent_dim=256): 409 | """ 410 | Parameters: 411 | input_dim (int) -- dimension of B 412 | output_dim (int) -- dimension of A 413 | norm_layer -- normalization layer 414 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 415 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 416 | latent_dim (int) -- the dimensionality of the latent space 417 | """ 418 | super(FcVed, self).__init__() 419 | 420 | # ENCODER 421 | # Layer 1 422 | self.encode_fc_1 = FCBlock(input_dim, dim_1B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 423 | activation=True) 424 | # Layer 2 425 | self.encode_fc_2 = FCBlock(dim_1B, dim_2B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 426 | activation=True) 427 | # Layer 3 428 | self.encode_fc_3 = FCBlock(dim_2B, dim_3_B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 429 | activation=True) 430 | # Layer 4 431 | self.encode_fc_mean = FCBlock(dim_3_B, latent_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, activation=False, 432 | normalization=False) 433 | self.encode_fc_log_var = FCBlock(dim_3_B, latent_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, activation=False, 434 | normalization=False) 435 | 436 | # DECODER 437 | # Layer 1 438 | self.decode_fc_z = FCBlock(latent_dim, dim_3_A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 439 | activation=True) 440 | # Layer 2 441 | self.decode_fc_2 = FCBlock(dim_3_A, dim_2A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 442 | activation=True) 443 | # Layer 3 444 | self.decode_fc_3 = FCBlock(dim_2A, dim_1A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 445 | activation=True) 446 | # Layer 4 447 | self.decode_fc_4 = FCBlock(dim_1A, output_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, activation=False, 448 | normalization=False) 449 | 450 | def encode(self, x): 451 | level_2 = self.encode_fc_1(x) 452 | level_3 = self.encode_fc_2(level_2) 453 | level_4 = self.encode_fc_3(level_3) 454 | latent_mean = self.encode_fc_mean(level_4) 455 | latent_log_var = self.encode_fc_log_var(level_4) 456 | return latent_mean, latent_log_var 457 | 458 | def reparameterize(self, mean, log_var): 459 | std = torch.exp(0.5 * log_var) 460 | eps = torch.randn_like(std) 461 | return eps.mul(std).add_(mean) 462 | 463 | def decode(self, z): 464 | level_1 = self.decode_fc_z(z) 465 | level_2 = self.decode_fc_2(level_1) 466 | level_3 = self.decode_fc_3(level_2) 467 | recon_A = self.decode_fc_4(level_3) 468 | return recon_A 469 | 470 | def forward(self, x): 471 | mean, log_var = self.encode(x) 472 | z = self.reparameterize(mean, log_var) 473 | recon_x = self.decode(z) 474 | return z, recon_x, mean, log_var 475 | 476 | 477 | class FcVedSep(nn.Module): 478 | """ 479 | Defines a fully-connected variational encoder-decoder model for DNA methylation to gene expression translation 480 | DNA methylation input separated by chromosome 481 | """ 482 | def __init__(self, input_dim, output_dim, norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0, dim_1B=128, dim_2B=1024, 483 | dim_3B=512, dim_1A=2048, dim_2A=1024, dim_3A=512, latent_dim=256): 484 | """ 485 | Construct a fully-connected variational autoencoder 486 | Parameters: 487 | input_dim (list) -- dimension list of B 488 | output_dim (int) -- dimension of A 489 | norm_layer -- normalization layer 490 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 491 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 492 | latent_dim (int) -- the dimensionality of the latent space 493 | """ 494 | super(FcVedSep, self).__init__() 495 | 496 | # ENCODER 497 | # Layer 1 498 | self.encode_fc_1_list = nn.ModuleList() 499 | for i in range(0, 23): 500 | self.encode_fc_1_list.append( 501 | FCBlock(input_dim[i], dim_1B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 502 | activation=True)) 503 | # Layer 2 504 | self.encode_fc_2 = FCBlock(dim_1B*23, dim_2B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 505 | activation=True) 506 | # Layer 3 507 | self.encode_fc_3 = FCBlock(dim_2B, dim_3B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 508 | activation=True) 509 | # Layer 4 510 | self.encode_fc_mean = FCBlock(dim_3B, latent_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, activation=False, 511 | normalization=False) 512 | self.encode_fc_log_var = FCBlock(dim_3B, latent_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, activation=False, 513 | normalization=False) 514 | 515 | # DECODER 516 | # Layer 1 517 | self.decode_fc_z = FCBlock(latent_dim, dim_3A, norm_layer=norm_layer, leaky_slope=leaky_slope, 518 | dropout_p=dropout_p, activation=True) 519 | # Layer 2 520 | self.decode_fc_2 = FCBlock(dim_3A, dim_2A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 521 | activation=True) 522 | # Layer 3 523 | self.decode_fc_3 = FCBlock(dim_2A, dim_1A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 524 | activation=True) 525 | # Layer 4 526 | self.decode_fc_4 = FCBlock(dim_1A, output_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, activation=False, 527 | normalization=False) 528 | 529 | def encode(self, x): 530 | level_2_list = [] 531 | for i in range(0, 23): 532 | level_2_list.append(self.encode_fc_1_list[i](x[i])) 533 | level_2 = torch.cat(level_2_list, 2) 534 | level_3 = self.encode_fc_2(level_2) 535 | level_4 = self.encode_fc_3(level_3) 536 | latent_mean = self.encode_fc_mean(level_4) 537 | latent_log_var = self.encode_fc_log_var(level_4) 538 | return latent_mean, latent_log_var 539 | 540 | def reparameterize(self, mean, log_var): 541 | std = torch.exp(0.5 * log_var) 542 | eps = torch.randn_like(std) 543 | return eps.mul(std).add_(mean) 544 | 545 | def decode(self, z): 546 | level_1 = self.decode_fc_z(z) 547 | level_2 = self.decode_fc_2(level_1) 548 | level_3 = self.decode_fc_3(level_2) 549 | recon_A = self.decode_fc_4(level_3) 550 | return recon_A 551 | 552 | def forward(self, x): 553 | mean, log_var = self.encode(x) 554 | z = self.reparameterize(mean, log_var) 555 | recon_A = self.decode(z) 556 | return z, recon_A, mean, log_var 557 | 558 | 559 | class FCG(nn.Module): 560 | """ 561 | Create a fully-connected generator network 562 | """ 563 | def __init__(self, input_dim, output_dim, norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0, latent_dim=256): 564 | """ 565 | Construct a one dimensional U-Net 566 | Parameters: 567 | input_dim (int) -- the dimension of the input omics data 568 | output_dim (int) -- the dimension of the output omics data 569 | norm_layer -- normalization layer 570 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 571 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 572 | latent_dim (int) -- the dimensionality of the latent space 573 | """ 574 | super(FCG, self).__init__() 575 | 576 | dim_1 = 4096 577 | dim_2 = 1024 578 | dim_3 = 512 579 | 580 | # dim_1 = 256 581 | # dim_2 = 128 582 | # dim_3 = 128 583 | 584 | mul_fc_block = [FCBlock(input_dim, dim_1, norm_layer, leaky_slope, dropout_p), 585 | FCBlock(dim_1, dim_2, norm_layer, leaky_slope, dropout_p), 586 | FCBlock(dim_2, dim_3, norm_layer, leaky_slope, dropout_p), 587 | FCBlock(dim_3, latent_dim, norm_layer, leaky_slope, dropout_p), 588 | FCBlock(latent_dim, dim_3, norm_layer, leaky_slope, dropout_p), 589 | FCBlock(dim_3, dim_2, norm_layer, leaky_slope, dropout_p), 590 | FCBlock(dim_2, dim_1, norm_layer, leaky_slope, dropout_p), 591 | nn.Linear(dim_1, output_dim) 592 | ] 593 | self.mul_fc = nn.Sequential(*mul_fc_block) 594 | 595 | def forward(self, x): 596 | return self.mul_fc(x) 597 | 598 | 599 | class FCGSep(nn.Module): 600 | """ 601 | Defines a fully-connected encoder-decoder model for DNA methylation to gene expression translation 602 | DNA methylation input separated by chromosome 603 | """ 604 | def __init__(self, input_dim, output_dim, norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0, dim_1B=256, dim_2B=1024, 605 | dim_3B=512, dim_1A=2048, dim_2A=1024, dim_3A=512, latent_dim=256): 606 | """ 607 | Construct a fully-connected variational autoencoder 608 | Parameters: 609 | input_dim (list) -- dimension list of B 610 | output_dim (int) -- dimension of A 611 | norm_layer -- normalization layer 612 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 613 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 614 | latent_dim (int) -- the dimensionality of the latent space 615 | """ 616 | super(FCGSep, self).__init__() 617 | 618 | # ENCODER 619 | # Layer 1 620 | self.encode_fc_1_list = nn.ModuleList() 621 | for i in range(0, 23): 622 | self.encode_fc_1_list.append( 623 | FCBlock(input_dim[i], dim_1B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 624 | activation=True)) 625 | # Layer 2 626 | self.encode_fc_2 = FCBlock(dim_1B*23, dim_2B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 627 | activation=True) 628 | # Layer 3 629 | self.encode_fc_3 = FCBlock(dim_2B, dim_3B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 630 | activation=True) 631 | # Layer 4 632 | self.encode_fc_4 = FCBlock(dim_3B, latent_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 633 | activation=True) 634 | # DECODER 635 | # Layer 1 636 | self.decode_fc_z = FCBlock(latent_dim, dim_3A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 637 | activation=True) 638 | # Layer 2 639 | self.decode_fc_2 = FCBlock(dim_3A, dim_2A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 640 | activation=True) 641 | # Layer 3 642 | self.decode_fc_3 = FCBlock(dim_2A, dim_1A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 643 | activation=True) 644 | # Layer 4 645 | self.decode_fc_4 = FCBlock(dim_1A, output_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 646 | activation=False, normalization=False) 647 | 648 | def encode(self, x): 649 | level_2_list = [] 650 | for i in range(0, 23): 651 | level_2_list.append(self.encode_fc_1_list[i](x[i])) 652 | level_2 = torch.cat(level_2_list, 2) 653 | level_3 = self.encode_fc_2(level_2) 654 | level_4 = self.encode_fc_3(level_3) 655 | latent = self.encode_fc_4(level_4) 656 | return latent 657 | 658 | def decode(self, z): 659 | level_1 = self.decode_fc_z(z) 660 | level_2 = self.decode_fc_2(level_1) 661 | level_3 = self.decode_fc_3(level_2) 662 | recon_A = self.decode_fc_4(level_3) 663 | return recon_A 664 | 665 | def forward(self, x): 666 | latent = self.encode(x) 667 | recon_A = self.decode(latent) 668 | return recon_A 669 | 670 | 671 | class FCGSingle(nn.Module): 672 | """ 673 | Defines a single hidden layer fully-connected model for DNA methylation to gene expression translation 674 | """ 675 | def __init__(self, input_dim, output_dim, latent_dim=4000): 676 | """ 677 | Construct a fully-connected variational autoencoder 678 | Parameters: 679 | input_dim (list) -- dimension list of B 680 | output_dim (int) -- dimension of A 681 | norm_layer -- normalization layer 682 | latent_dim (int) -- the dimensionality of the latent space 683 | """ 684 | super(FCGSingle, self).__init__() 685 | layers = [nn.Linear(input_dim, latent_dim), 686 | nn.Sigmoid(), 687 | nn.Linear(latent_dim, output_dim) 688 | ] 689 | self.layers = nn.Sequential(*layers) 690 | 691 | def forward(self, x): 692 | recon_A = self.layers(x) 693 | return recon_A 694 | 695 | 696 | class LinearRegression(nn.Module): 697 | """ 698 | Defines a linear regression model for DNA methylation to gene expression translation 699 | """ 700 | def __init__(self, input_dim, output_dim): 701 | """ 702 | Construct a linear regression 703 | Parameters: 704 | input_dim (list) -- dimension list of B 705 | output_dim (int) -- dimension of A 706 | """ 707 | super(LinearRegression, self).__init__() 708 | self.lr = nn.Linear(input_dim, output_dim) 709 | 710 | def forward(self, x): 711 | recon_A = self.lr(x) 712 | return recon_A 713 | 714 | 715 | class FCBlock(nn.Module): 716 | """ 717 | Linear => Norm1D => LeakyReLU 718 | """ 719 | def __init__(self, input_dim, output_dim, norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0, activation=True, normalization=True, activation_name='leakyrelu'): 720 | """ 721 | Construct a fully-connected block 722 | Parameters: 723 | input_dim (int) -- the dimension of the input tensor 724 | output_dim (int) -- the dimension of the output tensor 725 | norm_layer -- normalization layer 726 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 727 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 728 | activation (bool) -- need activation or not 729 | normalization (bool) -- need normalization or not 730 | activation_name (str) -- name of the activation function used in the FC block 731 | """ 732 | super(FCBlock, self).__init__() 733 | # Linear 734 | self.fc_block = [nn.Linear(input_dim, output_dim)] 735 | # Norm 736 | if normalization: 737 | # FC block doesn't support BatchNorm1d 738 | if isinstance(norm_layer, functools.partial) and norm_layer.func == nn.BatchNorm1d: 739 | norm_layer = nn.InstanceNorm1d 740 | self.fc_block.append(norm_layer(output_dim)) 741 | # Dropout 742 | if 0 < dropout_p <= 1: 743 | self.fc_block.append(nn.Dropout(p=dropout_p)) 744 | # LeakyReLU 745 | if activation: 746 | if activation_name.lower() == 'leakyrelu': 747 | self.fc_block.append(nn.LeakyReLU(negative_slope=leaky_slope, inplace=True)) 748 | elif activation_name.lower() == 'tanh': 749 | self.fc_block.append(nn.Tanh()) 750 | elif activation_name.lower() == 'sigmoid': 751 | self.fc_block.append(nn.Sigmoid()) 752 | else: 753 | raise NotImplementedError('Activation function [%s] is not implemented' % activation_name) 754 | 755 | self.fc_block = nn.Sequential(*self.fc_block) 756 | 757 | def forward(self, x): 758 | y = self.fc_block(x) 759 | return y 760 | 761 | 762 | class TransformerG(nn.Module): 763 | """ 764 | Create a transformer generator network 765 | """ 766 | def __init__(self, input_dim, output_dim, norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0, latent_dim=256, 767 | nhead=8): 768 | """ 769 | Construct a transformer generator 770 | Parameters: 771 | input_dim (int) -- the dimension of the input omics data 772 | output_dim (int) -- the dimension of the output omics data 773 | norm_layer -- normalization layer 774 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 775 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 776 | latent_dim (int) -- the dimensionality of the latent space 777 | nhead (int) -- the number of heads in the transformer encoder layer 778 | """ 779 | super(TransformerG, self).__init__() 780 | 781 | mul_fc_block = [FCBlock(input_dim, latent_dim, norm_layer, leaky_slope, dropout_p), 782 | nn.TransformerEncoderLayer(d_model=latent_dim, nhead=nhead), 783 | nn.Linear(latent_dim, output_dim) 784 | ] 785 | self.mul_fc = nn.Sequential(*mul_fc_block) 786 | 787 | def forward(self, x): 788 | return self.mul_fc(x) 789 | 790 | 791 | # DISCRIMINATOR 792 | 793 | # For one dimensional multi-layer convolution discriminator 794 | class MulLayerDiscriminator(nn.Module): 795 | """ 796 | Defines a one dimensional multi-layer convolution discriminator 797 | """ 798 | def __init__(self, input_1_chan_num, input_2_chan_num, filter_num=64, layer_num=3, norm_layer=nn.BatchNorm1d, 799 | leaky_slope=0.2): 800 | """ 801 | Construct a one dimensional multi-layer discriminator 802 | Parameters: 803 | input_1_chan_num (int) -- the number of channels in the first input omics data (A) 804 | input_2_chan_num (int) -- the number of channels in the second input omics data (B) 805 | filter_num (int) -- the number of filters in the first convolution layer 806 | layer_num (int) -- the number of convolution layers in the discriminator 807 | norm_layer -- normalization layer 808 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 809 | """ 810 | super(MulLayerDiscriminator, self).__init__() 811 | 812 | self.input_conv = DisInputConv(input_1_chan_num, input_2_chan_num, filter_num, leaky_slope) 813 | 814 | # create a list to store conv blocks 815 | mul_conv_block = [] 816 | conv_block_filter_num = filter_num * 2 817 | # the block number of the multi-layer convolution block should not exceed 6 818 | block_layer_num = min(layer_num, 6) 819 | for num in range(0, block_layer_num): 820 | # the filter number should not exceed 1024 821 | next_filter_num = min(conv_block_filter_num, 1024) 822 | mul_conv_block += [SingleConv1D(conv_block_filter_num, next_filter_num, 8, norm_layer, leaky_slope)] 823 | conv_block_filter_num = next_filter_num 824 | self.mul_conv = nn.Sequential(*mul_conv_block) 825 | 826 | # the output convolution layer of the discriminator 827 | self.output_conv = nn.Conv1d(conv_block_filter_num, 1, kernel_size=3, stride=1, padding=1) 828 | 829 | def forward(self, input1, input2): 830 | combined_input = self.input_conv(input1, input2) 831 | x = self.mul_conv(combined_input) 832 | return self.output_conv(x) 833 | 834 | 835 | class MulLayerDiscriminatorNew(nn.Module): 836 | """ 837 | Defines a one dimensional multi-layer convolution discriminator 838 | """ 839 | def __init__(self, input_1_chan_num, input_2_chan_num, a_dim, output_dim=64, filter_num=64, kernel_size=9, norm_layer=nn.BatchNorm1d, 840 | leaky_slope=0.2, dropout_p=0): 841 | """ 842 | Construct a one dimensional multi-layer discriminator 843 | Parameters: 844 | input_1_chan_num (int) -- the number of channels in the first input omics data (A) 845 | input_2_chan_num (int) -- the number of channels in the second input omics data (B) 846 | filter_num (int) -- the number of filters in the first convolution layer 847 | layer_num (int) -- the number of convolution layers in the discriminator 848 | norm_layer -- normalization layer 849 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 850 | """ 851 | super(MulLayerDiscriminatorNew, self).__init__() 852 | 853 | self.input_conv = DisInputConvUp(input_1_chan_num, input_2_chan_num, filter_num, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope) 854 | 855 | # create a list to store conv blocks 856 | self.mul_conv = nn.Sequential( 857 | # 128 -> 256 858 | DownSample(filter_num * 2, filter_num * 4, down_ratio=4, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p), 859 | # 256 -> 512 860 | DownSample(filter_num * 4, filter_num * 8, down_ratio=4, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p), 861 | # 512 -> 1024 862 | DownSample(filter_num * 8, filter_num * 16, down_ratio=4, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p), 863 | # 1024 -> 512 864 | SingleConv1D(filter_num * 16, filter_num * 8, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope), 865 | # 512 -> 256 866 | SingleConv1D(filter_num * 8, filter_num * 4, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope), 867 | # 256 -> 128 868 | SingleConv1D(filter_num * 4, filter_num * 2, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope), 869 | # 128 -> 64 870 | SingleConv1D(filter_num * 2, filter_num, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope), 871 | # 64 -> 1 872 | OutputConv(filter_num, 1), 873 | # FC 874 | FCBlock(a_dim // 4 // 4 // 4, output_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p), 875 | nn.Linear(output_dim, output_dim) 876 | ) 877 | 878 | def forward(self, input1, input2): 879 | combined_input = self.input_conv(input1, input2) 880 | x = self.mul_conv(combined_input) 881 | return x 882 | 883 | 884 | class MulLayerDiscriminatorNewS(nn.Module): 885 | """ 886 | Defines a one dimensional multi-layer convolution discriminator 887 | """ 888 | def __init__(self, input_1_chan_num, input_2_chan_num, filter_num=64, kernel_size=9, norm_layer=nn.BatchNorm1d, 889 | leaky_slope=0.2, dropout_p=0): 890 | """ 891 | Construct a one dimensional multi-layer discriminator 892 | Parameters: 893 | input_1_chan_num (int) -- the number of channels in the first input omics data (A) 894 | input_2_chan_num (int) -- the number of channels in the second input omics data (B) 895 | filter_num (int) -- the number of filters in the first convolution layer 896 | layer_num (int) -- the number of convolution layers in the discriminator 897 | norm_layer -- normalization layer 898 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 899 | """ 900 | super(MulLayerDiscriminatorNewS, self).__init__() 901 | 902 | self.input_conv = DisInputConvUp(input_1_chan_num, input_2_chan_num, filter_num, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope) 903 | 904 | # create a list to store conv blocks 905 | self.mul_conv = nn.Sequential( 906 | # 128 -> 256 907 | DownSample(filter_num * 2, filter_num * 4, down_ratio=4, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p), 908 | # 256 -> 512 909 | DownSample(filter_num * 4, filter_num * 8, down_ratio=4, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p), 910 | # 512 -> 1024 911 | DownSample(filter_num * 8, filter_num * 16, down_ratio=4, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p), 912 | # 1024 -> 1 913 | nn.Conv1d(filter_num * 16, 1, kernel_size=3, stride=1, padding=1), 914 | ) 915 | 916 | def forward(self, input1, input2): 917 | combined_input = self.input_conv(input1, input2) 918 | x = self.mul_conv(combined_input) 919 | return x 920 | 921 | 922 | class DisInputConv(nn.Module): 923 | """ 924 | The input convolution block for the conditional GAN multi-layer discriminator 925 | The input of this block are the two different omics data 926 | """ 927 | def __init__(self, input_1_chan_num, input_2_chan_num, output_chan_num, leaky_slope=0.2): 928 | """ 929 | Construct the input convolution layer 930 | Parameters: 931 | input_1_chan_num (int) -- the number of channels of the first input omics data (A) 932 | input_2_chan_num (int) -- the number of channels of the second input omics data (B) 933 | output_chan_num (int) -- the number of channels of the output tensor 934 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 935 | """ 936 | super(DisInputConv, self).__init__() 937 | 938 | self.input_1_conv_layer = nn.Sequential( 939 | nn.Conv1d(input_1_chan_num, output_chan_num, kernel_size=9, stride=1, padding=4), 940 | nn.LeakyReLU(negative_slope=leaky_slope, inplace=True) 941 | ) 942 | self.input_2_conv_layer = nn.Sequential( 943 | nn.Conv1d(input_2_chan_num, output_chan_num, kernel_size=9, stride=7, padding=1), 944 | nn.LeakyReLU(negative_slope=leaky_slope, inplace=True) 945 | ) 946 | 947 | def forward(self, input1, input2): 948 | x1 = self.input_1_conv_layer(input1) 949 | x2 = self.input_2_conv_layer(input2) 950 | x2_crop = x2[:, :, 0:x1.shape[2]] 951 | # concat the two input omics data together 952 | x = torch.cat([x1, x2_crop], dim=1) 953 | return x 954 | 955 | 956 | class DisInputConvUp(nn.Module): 957 | """ 958 | The input convolution block for the conditional GAN multi-layer discriminator 959 | The input of this block are the two different omics data 960 | """ 961 | def __init__(self, input_1_chan_num, input_2_chan_num, output_chan_num, up_ratio=2, kernel_size=9, norm_layer=nn.BatchNorm1d, leaky_slope=0.2): 962 | """ 963 | Construct the input convolution layer 964 | Parameters: 965 | input_1_chan_num (int) -- the number of channels of the first input omics data (A) 966 | input_2_chan_num (int) -- the number of channels of the second input omics data (B) 967 | output_chan_num (int) -- the number of channels of the output tensor 968 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 969 | """ 970 | super(DisInputConvUp, self).__init__() 971 | 972 | # Only if the norm method is instance norm we use bias for the corresponding conv layer 973 | if type(norm_layer) == functools.partial: 974 | use_bias = norm_layer.func == nn.InstanceNorm1d 975 | else: 976 | use_bias = norm_layer == nn.InstanceNorm1d 977 | 978 | self.input_1_conv = DoubleConv1D(input_1_chan_num, output_chan_num, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope) 979 | self.input_2_conv = DoubleConv1D(input_2_chan_num, output_chan_num, kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope) 980 | 981 | self.input_1_out = nn.Conv1d(output_chan_num, output_chan_num, kernel_size=kernel_size, padding=kernel_size // 2, bias=use_bias) 982 | 983 | self.input_2_out = nn.ConvTranspose1d(output_chan_num, output_chan_num, kernel_size=up_ratio, stride=up_ratio) 984 | 985 | def forward(self, input1, input2): 986 | x1 = self.input_1_conv(input1) 987 | x1_out = self.input_1_out(x1) 988 | x2 = self.input_2_conv(input2) 989 | x2_out = self.input_2_out(x2) 990 | x2_crop = x2_out[:, :, 0:x1.shape[2]] 991 | # concat the two input omics data together 992 | x = torch.cat([x1_out, x2_crop], dim=1) 993 | return x 994 | 995 | 996 | class MulLayerDiscriminatorSep(nn.Module): 997 | """ 998 | Defines a one dimensional multi-layer convolution discriminator 999 | """ 1000 | def __init__(self, input_1_chan_num, input_2_chan_num, filter_num=64, layer_num=3, norm_layer=nn.BatchNorm1d, 1001 | leaky_slope=0.2): 1002 | """ 1003 | Construct a one dimensional multi-layer discriminator 1004 | Parameters: 1005 | input_1_chan_num (int) -- the number of channels in the first input omics data (A) 1006 | input_2_chan_num (int) -- the number of channels in the second input omics data (B) 1007 | filter_num (int) -- the number of filters in the first convolution layer 1008 | layer_num (int) -- the number of convolution layers in the discriminator 1009 | norm_layer -- normalization layer 1010 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 1011 | """ 1012 | super(MulLayerDiscriminatorSep, self).__init__() 1013 | 1014 | self.input_conv = DisInputConvSep(input_1_chan_num, input_2_chan_num, filter_num, leaky_slope) 1015 | 1016 | # create a list to store conv blocks 1017 | mul_conv_block = [] 1018 | conv_block_filter_num = filter_num * 2 1019 | # the block number of the multi-layer convolution block should not exceed 6 1020 | block_layer_num = min(layer_num, 6) 1021 | for num in range(0, block_layer_num): 1022 | # the filter number should not exceed 1024 1023 | next_filter_num = min(conv_block_filter_num, 1024) 1024 | mul_conv_block += [SingleConv1D(conv_block_filter_num, next_filter_num, 8, norm_layer, leaky_slope)] 1025 | conv_block_filter_num = next_filter_num 1026 | self.mul_conv = nn.Sequential(*mul_conv_block) 1027 | 1028 | # the output convolution layer of the discriminator 1029 | self.output_conv = nn.Conv1d(conv_block_filter_num, 1, kernel_size=3, stride=1, padding=1) 1030 | 1031 | def forward(self, input1, input2): 1032 | combined_input = self.input_conv(input1, input2) 1033 | x = self.mul_conv(combined_input) 1034 | return self.output_conv(x) 1035 | 1036 | 1037 | class DisInputConvSep(nn.Module): 1038 | """ 1039 | The input convolution block for the conditional GAN multi-layer discriminator 1040 | The input of this block are the two different omics data 1041 | """ 1042 | def __init__(self, input_1_chan_num, input_2_chan_num, output_chan_num, leaky_slope=0.2): 1043 | """ 1044 | Construct the input convolution layer 1045 | Parameters: 1046 | input_1_chan_num (int) -- the number of channels of the first input omics data (A) 1047 | input_2_chan_num (int) -- the number of channels of the second input omics data (B) 1048 | output_chan_num (int) -- the number of channels of the output tensor 1049 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 1050 | """ 1051 | super(DisInputConvSep, self).__init__() 1052 | 1053 | self.input_1_conv_layer = nn.Sequential( 1054 | nn.Conv1d(input_1_chan_num, output_chan_num, kernel_size=9, stride=1, padding=4), 1055 | nn.LeakyReLU(negative_slope=leaky_slope, inplace=True) 1056 | ) 1057 | self.input_2_conv_layer = nn.Sequential( 1058 | nn.Conv1d(input_2_chan_num, output_chan_num, kernel_size=9, stride=7, padding=1), 1059 | nn.LeakyReLU(negative_slope=leaky_slope, inplace=True) 1060 | ) 1061 | 1062 | def forward(self, input1, input2): 1063 | input2_list = [] 1064 | for i in range(0, 23): 1065 | input2_list.append(input2[i]) 1066 | input2_cat = torch.cat(input2_list, 2) 1067 | 1068 | x1 = self.input_1_conv_layer(input1) 1069 | x2 = self.input_2_conv_layer(input2_cat) 1070 | x2_crop = x2[:, :, 0:x1.shape[2]] 1071 | # concat the two input omics data together 1072 | x = torch.cat([x1, x2_crop], dim=1) 1073 | return x 1074 | 1075 | 1076 | # For fully-connected discriminator 1077 | class FCD(nn.Module): 1078 | """ 1079 | Defines a fully-connected discriminator 1080 | """ 1081 | def __init__(self, input_1_dim, input_2_dim, norm_layer=nn.BatchNorm1d, leaky_slope=0.2, output_dim=16): 1082 | """ 1083 | Construct a fully-connected discriminator 1084 | Parameters: 1085 | input_1_dim (int) -- the dimension of the first input omics data (A) 1086 | input_2_dim (int) -- the dimension of the second input omics data (B) 1087 | norm_layer -- normalization layer 1088 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 1089 | output_dim (int) -- the output dimension of the discriminator 1090 | """ 1091 | super(FCD, self).__init__() 1092 | 1093 | dim_1 = 128 1094 | dim_2 = 64 1095 | dim_3 = 32 1096 | 1097 | mul_fc_block = [FCBlock(input_1_dim+input_2_dim, dim_1, norm_layer, leaky_slope), 1098 | FCBlock(dim_1, dim_2, norm_layer, leaky_slope), 1099 | FCBlock(dim_2, dim_3, norm_layer, leaky_slope), 1100 | nn.Linear(dim_3, output_dim)] 1101 | self.mul_fc = nn.Sequential(*mul_fc_block) 1102 | 1103 | def forward(self, input1, input2): 1104 | combined_input = torch.cat([input1, input2], dim=2) 1105 | output = self.mul_fc(combined_input) 1106 | return output 1107 | 1108 | 1109 | class FCDSingle(nn.Module): 1110 | """ 1111 | Defines a single hidden layer fully-connected discriminator 1112 | """ 1113 | def __init__(self, input_1_dim, input_2_dim, hidden_dim=64, output_dim=16): 1114 | """ 1115 | Construct a fully-connected discriminator 1116 | Parameters: 1117 | input_1_dim (int) -- the dimension of the first input omics data (A) 1118 | input_2_dim (int) -- the dimension of the second input omics data (B) 1119 | output_dim (int) -- the output dimension of the discriminator 1120 | """ 1121 | super(FCDSingle, self).__init__() 1122 | 1123 | layers = [nn.Linear(input_1_dim+input_2_dim, hidden_dim), 1124 | nn.LeakyReLU(negative_slope=0.2, inplace=True), 1125 | nn.Linear(hidden_dim, output_dim) 1126 | ] 1127 | self.layers = nn.Sequential(*layers) 1128 | 1129 | def forward(self, input1, input2): 1130 | combined_input = torch.cat([input1, input2], dim=2) 1131 | output = self.layers(combined_input) 1132 | return output 1133 | 1134 | 1135 | class FCDSep(nn.Module): 1136 | """ 1137 | Defines a fully-connected discriminator 1138 | DNA methylation input separated by chromosome 1139 | """ 1140 | def __init__(self, input_1_dim, input_2_dim, norm_layer=nn.BatchNorm1d, leaky_slope=0.2, dropout_p=0, dim_1B=16, 1141 | dim_2B=64, dim_1A=128, dim_2A=64, dim_3=32, output_dim=16): 1142 | """ 1143 | Construct a fully-connected discriminator 1144 | Parameters: 1145 | input_1_dim (int) -- the dimension of the first input omics data (A) 1146 | input_2_dim (int) -- the dimension of the second input omics data (B) 1147 | norm_layer -- normalization layer 1148 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 1149 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 1150 | output_dim (int) -- the output dimension of the discriminator 1151 | """ 1152 | super(FCDSep, self).__init__() 1153 | 1154 | A_dim = input_1_dim 1155 | B_dim_list = input_2_dim 1156 | 1157 | # ENCODER 1158 | # Layer 1 1159 | self.encode_fc_1B_list = nn.ModuleList() 1160 | for i in range(0, 23): 1161 | self.encode_fc_1B_list.append( 1162 | FCBlock(B_dim_list[i], dim_1B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 1163 | activation=True)) 1164 | self.encode_fc_1A = FCBlock(A_dim, dim_1A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 1165 | activation=True) 1166 | # Layer 2 1167 | self.encode_fc_2B = FCBlock(dim_1B * 23, dim_2B, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 1168 | activation=True) 1169 | self.encode_fc_2A = FCBlock(dim_1A, dim_2A, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 1170 | activation=True) 1171 | # Layer 3 1172 | self.encode_fc_3 = FCBlock(dim_2B + dim_2A, dim_3, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 1173 | activation=True) 1174 | # Layer 4 1175 | self.encode_fc_4 = FCBlock(dim_3, output_dim, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p, 1176 | activation=False) 1177 | 1178 | def forward(self, input1, input2): 1179 | level_2_B_list = [] 1180 | for i in range(0, 23): 1181 | level_2_B_list.append(self.encode_fc_1B_list[i](input2[i])) 1182 | level_2_B = torch.cat(level_2_B_list, 2) 1183 | level_2_A = self.encode_fc_1A(input1) 1184 | level_3_B = self.encode_fc_2B(level_2_B) 1185 | level_3_A = self.encode_fc_2A(level_2_A) 1186 | level_3 = torch.cat((level_3_B, level_3_A), 2) 1187 | level_4 = self.encode_fc_3(level_3) 1188 | output = self.encode_fc_4(level_4) 1189 | return output 1190 | 1191 | 1192 | class TransformerD(nn.Module): 1193 | """ 1194 | Create a transformer discriminator network 1195 | """ 1196 | def __init__(self, input_1_dim, input_2_dim, hidden_dim=64, output_dim=16, nhead=8): 1197 | """ 1198 | Construct a transformer discriminator 1199 | Parameters: 1200 | input_1_dim (int) -- the dimension of the first input omics data (A) 1201 | input_2_dim (int) -- the dimension of the second input omics data (B) 1202 | output_dim (int) -- the output dimension of the discriminator 1203 | nhead (int) -- the number of heads in the transformer encoder layer 1204 | """ 1205 | super(TransformerD, self).__init__() 1206 | 1207 | layers = [nn.Linear(input_1_dim+input_2_dim, hidden_dim), 1208 | nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=nhead), 1209 | nn.LeakyReLU(negative_slope=0.2, inplace=True), 1210 | nn.Linear(hidden_dim, output_dim) 1211 | ] 1212 | self.layers = nn.Sequential(*layers) 1213 | 1214 | def forward(self, input1, input2): 1215 | combined_input = torch.cat([input1, input2], dim=2) 1216 | output = self.layers(combined_input) 1217 | return output 1218 | 1219 | 1220 | def init_weights(net, init_type='normal', init_gain=0.02): 1221 | """ 1222 | Initialize network weights. 1223 | Parameters: 1224 | net (nn.Module) -- the network to be initialized 1225 | init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal 1226 | init_gain (float) -- scaling factor for normal, xavier and orthogonal. 1227 | """ 1228 | # define the initialization function 1229 | def init_func(m): 1230 | classname = m.__class__.__name__ 1231 | if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): 1232 | if init_type == 'normal': 1233 | init.normal_(m.weight.data, 0.0, init_gain) 1234 | elif init_type == 'xavier': 1235 | init.xavier_normal_(m.weight.data, gain=init_gain) 1236 | elif init_type == 'kaiming': 1237 | init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') 1238 | elif init_type == 'orthogonal': 1239 | init.orthogonal_(m.weight.data, gain=init_gain) 1240 | else: 1241 | raise NotImplementedError('initialization method [%s] is not implemented' % init_type) 1242 | if hasattr(m, 'bias') and m.bias is not None: 1243 | init.constant_(m.bias.data, 0.0) 1244 | elif classname.find('BatchNorm1d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies. 1245 | init.normal_(m.weight.data, 1.0, init_gain) 1246 | init.constant_(m.bias.data, 0.0) 1247 | 1248 | print('Initialize network with %s' % init_type) 1249 | net.apply(init_func) # apply the initialization function 1250 | 1251 | 1252 | def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]): 1253 | """ 1254 | Initialize a network: 1255 | 1. register CPU/GPU device (with multi-GPU support); 1256 | 2. initialize the network weights 1257 | Parameters: 1258 | net (nn.Module) -- the network to be initialized 1259 | init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal 1260 | init_gain (float) -- scaling factor for normal, xavier and orthogonal. 1261 | gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 1262 | Return an initialized network. 1263 | """ 1264 | if len(gpu_ids) > 0: 1265 | assert(torch.cuda.is_available()) 1266 | net.to(gpu_ids[0]) 1267 | # multi-GPUs 1268 | net = torch.nn.DataParallel(net, gpu_ids) 1269 | init_weights(net, init_type, init_gain=init_gain) 1270 | return net 1271 | 1272 | 1273 | def get_norm_layer(norm_type='batch'): 1274 | """ 1275 | Return a normalization layer 1276 | Parameters: 1277 | norm_type (str) -- the type of normalization applied to the model, default to use batch normalization, options: [batch | instance | none ] 1278 | """ 1279 | if norm_type == 'batch': 1280 | norm_layer = functools.partial(nn.BatchNorm1d, affine=True, track_running_stats=True) 1281 | elif norm_type == 'instance': 1282 | norm_layer = functools.partial(nn.InstanceNorm1d, affine=False, track_running_stats=False) 1283 | elif norm_type == 'none': 1284 | norm_layer = lambda x: Identity() 1285 | else: 1286 | raise NotImplementedError('normalization method [%s] is not found' % norm_type) 1287 | return norm_layer 1288 | 1289 | 1290 | def define_G(input_chan_num, output_chan_num, netG, A_dim, B_dim, gen_filter_num=64, kernel_size=9, norm_type='batch', 1291 | init_type='normal', init_gain=0.02, gpu_ids=[], leaky_slope=0.2, dropout_p=0, latent_dim=256): 1292 | """ 1293 | Create a generator 1294 | 1295 | Parameters: 1296 | input_chan_num (int) -- the number of channels in input omics data 1297 | output_chan_num (int) -- the number of channels in output omics data 1298 | netG (str) -- the name of the generator architecture, default: unet_1d 1299 | A_dim (int) -- the dimension of omics type A 1300 | B_dim (int) -- the dimension of omics type B 1301 | gen_filter_num (int) -- the number of filters in the first convolution layer in the generator 1302 | kernel_size (int) -- the kernel size of convolution layers 1303 | norm_type (str) -- the name of normalization layers used in the network, default: batch 1304 | init_type (str) -- the name of our initialization method 1305 | init_gain (float) -- scaling factor for normal, xavier and orthogonal initialization methods 1306 | gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1 1307 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 1308 | dropout_p (float) -- probability of an element to be zeroed in a dropout layer 1309 | latent_dim (int) -- latent dimension for the encoder-decoder model 1310 | 1311 | Returns a generator 1312 | 1313 | The default implementation of the generator is the one dimensional U-Net architecture. 1314 | 1315 | This architecture is modified from the original architecture mentioned in the U-Net paper: https://link.springer.com/chapter/10.1007/978-3-319-24574-4_28 1316 | 1317 | The generator has been initialized by . 1318 | """ 1319 | 1320 | net = None 1321 | 1322 | # get the normalization layer 1323 | norm_layer = get_norm_layer(norm_type=norm_type) 1324 | 1325 | if netG == 'unet_de': 1326 | net = UNetDe(input_chan_num, output_chan_num, output_dim=A_dim, filter_num=gen_filter_num, 1327 | kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 1328 | elif netG == 'unet_in': 1329 | net = UNetIn(input_chan_num, output_chan_num, output_dim=A_dim, filter_num=gen_filter_num, 1330 | kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 1331 | elif netG == 'unet_s_de': 1332 | net = UNetSDe(input_chan_num, output_chan_num, output_dim=A_dim, filter_num=gen_filter_num, 1333 | kernel_size=kernel_size, norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 1334 | elif netG == 'fcg': 1335 | net = FCG(B_dim, A_dim, norm_layer, leaky_slope, dropout_p, latent_dim=latent_dim) 1336 | elif netG == 'fcg_sep': 1337 | net = FCGSep(B_dim, A_dim, norm_layer, leaky_slope, dropout_p, latent_dim=latent_dim) 1338 | elif netG == 'fcg_single': 1339 | net = FCGSingle(B_dim, A_dim) 1340 | elif netG == 'fc_ved': 1341 | net = FcVed(B_dim, A_dim, norm_layer, leaky_slope, dropout_p, latent_dim=latent_dim) 1342 | elif netG == 'fc_ved_sep': 1343 | net = FcVedSep(B_dim, A_dim, norm_layer, leaky_slope, dropout_p, latent_dim=latent_dim) 1344 | elif netG == 'linear_regression': 1345 | net = LinearRegression(B_dim, A_dim) 1346 | elif netG == 'transformer_g': 1347 | net = TransformerG(B_dim, A_dim, norm_layer, leaky_slope, dropout_p, latent_dim=latent_dim) 1348 | else: 1349 | raise NotImplementedError('Generator model name [%s] is not recognized' % netG) 1350 | 1351 | return init_net(net, init_type, init_gain, gpu_ids) 1352 | 1353 | 1354 | def define_D(input_1_chan_num, input_2_chan_num, dis_filter_num, netD, A_dim, B_dim, layer_num_D=3, norm_type='batch', 1355 | init_type='normal', init_gain=0.02, gpu_ids=[], leaky_slope=0.2, dropout_p=0): 1356 | """ 1357 | Create a discriminator 1358 | 1359 | Parameters: 1360 | input_1_chan_num (int) -- the number of channels in the first input omics data 1361 | input_2_chan_num (int) -- the number of channels in the second input omics data 1362 | dis_filter_num (int) -- the number of filters in the first convolution layer in the discriminator 1363 | netD (str) -- the name of the discriminator architecture, default: patch_gan_1d 1364 | A_dim (int) -- the dimension of omics type A 1365 | B_dim (int) -- the dimension of omics type B 1366 | layer_num_D (int) -- the number of convolution layers in the discriminator 1367 | norm_type (str) -- the type of normalization layers used in the network 1368 | init_type (str) -- the name of the initialization method 1369 | init_gain (float) -- scaling factor for normal, xavier and orthogonal initialization methods 1370 | gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1 1371 | leaky_slope (float) -- the negative slope of the Leaky ReLU activation function 1372 | output_dim (int) -- the output dimension of the discriminator 1373 | 1374 | Returns a discriminator 1375 | 1376 | The discriminator has been initialized by . 1377 | """ 1378 | net = None 1379 | norm_layer = get_norm_layer(norm_type=norm_type) 1380 | 1381 | if netD == 'multi_conv': # default one dimensional multi-layer convolution discriminator 1382 | net = MulLayerDiscriminator(input_1_chan_num, input_2_chan_num, filter_num=dis_filter_num, 1383 | layer_num=layer_num_D, norm_layer=norm_layer, leaky_slope=leaky_slope) 1384 | elif netD == 'multi_conv_sep': # default one dimensional multi-layer convolution discriminator 1385 | net = MulLayerDiscriminatorSep(input_1_chan_num, input_2_chan_num, filter_num=dis_filter_num, 1386 | layer_num=layer_num_D, norm_layer=norm_layer, leaky_slope=leaky_slope) 1387 | elif netD == 'multi_conv_new': # default one dimensional multi-layer convolution discriminator 1388 | net = MulLayerDiscriminatorNew(input_1_chan_num, input_2_chan_num, a_dim=A_dim, filter_num=dis_filter_num, 1389 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 1390 | elif netD == 'multi_conv_new_s': # default one dimensional multi-layer convolution discriminator 1391 | net = MulLayerDiscriminatorNewS(input_1_chan_num, input_2_chan_num, filter_num=dis_filter_num, 1392 | norm_layer=norm_layer, leaky_slope=leaky_slope, dropout_p=dropout_p) 1393 | elif netD == 'fcd': 1394 | net = FCD(A_dim, B_dim, norm_layer, leaky_slope) 1395 | elif netD == 'fcd_sep': 1396 | net = FCDSep(A_dim, B_dim, norm_layer, leaky_slope) 1397 | elif netD == 'fcd_single': 1398 | net = FCDSingle(A_dim, B_dim) 1399 | elif netD == 'transformer_d': 1400 | net = TransformerD(A_dim, B_dim) 1401 | else: 1402 | raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD) 1403 | return init_net(net, init_type, init_gain, gpu_ids) 1404 | 1405 | 1406 | def get_scheduler(optimizer, param): 1407 | """ 1408 | Return a learning rate scheduler 1409 | 1410 | Parameters: 1411 | optimizer (opt class) -- the optimizer of the network 1412 | param (params class) -- param.lr_policy is the name of learning rate policy: linear | step | plateau | cosine 1413 | 1414 | For 'linear', we keep the same learning rate for the first epochs and linearly decay the rate to zero 1415 | over the next epochs. 1416 | 1417 | """ 1418 | if param.lr_policy == 'linear': 1419 | def lambda_rule(epoch): 1420 | lr_lambda = 1.0 - max(0, epoch + param.epoch_count - param.epoch_num) / float(param.epoch_num_decay + 1) 1421 | return lr_lambda 1422 | # lr_scheduler is imported from torch.optim 1423 | scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) 1424 | elif param.lr_policy == 'step': 1425 | scheduler = lr_scheduler.StepLR(optimizer, step_size=param.decay_step_size, gamma=0.1) 1426 | elif param.lr_policy == 'plateau': 1427 | scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5) 1428 | elif param.lr_policy == 'cosine': 1429 | scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=param.epoch_num, eta_min=0) 1430 | else: 1431 | return NotImplementedError('Learning rate policy [%s] is not found', param.lr_policy) 1432 | return scheduler 1433 | -------------------------------------------------------------------------------- /params/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | This package deal with the parameters of the project, includes train params, test params and their father class basic 3 | params. 4 | """ -------------------------------------------------------------------------------- /params/basic_params.py: -------------------------------------------------------------------------------- 1 | import time 2 | import argparse 3 | import torch 4 | import os 5 | import models 6 | from util import util 7 | 8 | 9 | class BasicParams: 10 | """ 11 | This class is the father class of both TrainParams and TestParams. This class define the parameters used in both 12 | training and testing. 13 | """ 14 | 15 | def __init__(self): 16 | """ 17 | Reset the class. Indicates the class hasn't been initialized 18 | """ 19 | self.initialized = False 20 | self.isTrain = True 21 | self.isTest = True 22 | 23 | def initialize(self, parser): 24 | """ 25 | Define the common parameters used in both training and testing. 26 | """ 27 | parser.add_argument('--gpu_ids', type=str, default='0', 28 | help='which GPU would like to use: e.g. 0 or 0,1. -1 for CPU') 29 | parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', 30 | help='models, settings and intermediate results are saved in folder in this directory') 31 | parser.add_argument('--experiment_name', type=str, default='test', 32 | help='name of the folder in the checkpoint directory') 33 | # Dataset parameters 34 | parser.add_argument('--data_root', type=str, default='./data', 35 | help='path to input data') 36 | parser.add_argument('--dataset_mode', type=str, default='single_file', 37 | help='choose the dataset mode, options: [separate_files | single_file]') 38 | parser.add_argument('--batch_size', type=int, default=4, 39 | help='input data batch size') 40 | parser.add_argument('--num_threads', default=6, type=int, 41 | help='number of threads for loading data') 42 | parser.add_argument('--set_pin_memory', action='store_true', 43 | help='set pin_memory in the dataloader to increase data loading performance') 44 | parser.add_argument('--file_format', type=str, default='tsv', 45 | help='file format of the omics data, options: [tsv | csv | hdf]') 46 | parser.add_argument('--use_sample_list', action='store_true', 47 | help='provide a subset sample list of the dataset, store in the path data_root/sample_list.tsv, if False use all the samples') 48 | parser.add_argument('--use_feature_lists', action='store_true', 49 | help='provide feature lists of the input omics data, e.g. data_root/feature_list_A.tsv, if False use all the features') 50 | parser.add_argument('--detect_na', action='store_true', 51 | help='detect missing value markers during data loading, stay False can improve the loading performance') 52 | parser.add_argument('--zo_norm', action='store_true', 53 | help='built-in 0-1 normalisation') 54 | 55 | # Model parameters 56 | parser.add_argument('--model', type=str, default='c_gan', 57 | help='chooses which model want to use, options: [c_gan | c_gan_ved]') 58 | parser.add_argument('--norm_type', type=str, default='batch', 59 | help='the type of normalization applied to the model, default to use batch normalization, options: [batch | instance | none ]') 60 | parser.add_argument('--netG', type=str, default='fcg', 61 | help='specify generator architecture, default is the one dimensional U-Net architecture, options: [fc_ved | fc_ved_sep | fcg | fcg_sep | fcg_single | unet_de | unet_s_de | unet_in | linear_regression]') 62 | parser.add_argument('--netD', type=str, default='fcd', 63 | help='specify discriminator architecture, default is the one dimensional multi-layer convolution discriminator, options: [ fcd | fcd_sep | multi_conv | multi_conv_new ]') 64 | parser.add_argument('--input_chan_num', type=int, default=1, 65 | help='number of input omics data channels, default is 1, if there are different measurements they can be different channels') 66 | parser.add_argument('--output_chan_num', type=int, default=1, 67 | help='number of output omics channels, default is 1') 68 | parser.add_argument('--gen_filter_num', type=int, default=64, 69 | help='number of filters in the last convolution layer in the generator') 70 | parser.add_argument('--dis_filter_num', type=int, default=64, 71 | help='number of filters in the last convolution layer in the discriminator') 72 | parser.add_argument('--layer_num_D', type=int, default=3, 73 | help='the number of convolution layer in the discriminator') 74 | parser.add_argument('--conv_k_size', type=int, default=9, 75 | help='the kernel size of convolution layer, default kernel size is 9, the kernel is one dimensional.') 76 | parser.add_argument('--dropout_p', type=float, default=0, 77 | help='probability of an element to be zeroed in a dropout layer, default is 0 which means no dropout.') 78 | parser.add_argument('--leaky_slope', type=float, default=0.2, 79 | help='the negative slope of the Leaky ReLU activation function') 80 | parser.add_argument('--init_type', type=str, default='normal', 81 | help='choose the method of network initialization, options: [normal | xavier | kaiming | orthogonal]') 82 | parser.add_argument('--init_gain', type=float, default=0.02, 83 | help='scaling factor for normal, xavier and orthogonal initialization methods') 84 | parser.add_argument('--deterministic', action='store_true', 85 | help='make the model deterministic for reproduction if set true') 86 | parser.add_argument('--seed', type=int, default=42, 87 | help='random seed') 88 | parser.add_argument('--dist_loss', type=str, default='L1', 89 | help='chooses the distance loss between the generated value and the ground truth, options: [BCE | MSE | L1]') 90 | 91 | # Additional parameters 92 | parser.add_argument('--detail', action='store_true', 93 | help='print more detailed information if set true') 94 | parser.add_argument('--epoch_to_load', type=str, default='latest', 95 | help='The epoch number to load, set latest to load latest cached model') 96 | parser.add_argument('--experiment_to_load', type=str, default='test', 97 | help='the experiment to load') 98 | 99 | self.initialized = True # set the initialized to True after we define the parameters of the project 100 | return parser 101 | 102 | def get_params(self): 103 | """ 104 | Initialize our parser with basic parameters once. 105 | Add additional model-specific parameters. 106 | """ 107 | if not self.initialized: # check if this object has been initialized 108 | # if not create a new parser object 109 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) 110 | # use our method to initialize the parser with the predefined arguments 111 | parser = self.initialize(parser) 112 | 113 | # get the basic parameters 114 | param, _ = parser.parse_known_args() 115 | 116 | # modify model-related parser options 117 | model_name = param.model 118 | model_param_setter = models.get_param_setter(model_name) 119 | parser = model_param_setter(parser, self.isTrain) 120 | 121 | # save and return the parser 122 | self.parser = parser 123 | return parser.parse_args() 124 | 125 | def print_params(self, param): 126 | """ 127 | Print welcome words and command line parameters. 128 | Save the command line parameters in a txt file to the disk 129 | """ 130 | message = '' 131 | message += '\nWelcome to OmiTrans\nby Xiaoyu Zhang x.zhang18@imperial.ac.uk\n\n' 132 | message += '-----------------------Running Parameters-----------------------\n' 133 | for key, value in sorted(vars(param).items()): 134 | comment = '' 135 | default = self.parser.get_default(key) 136 | if value != default: 137 | comment = '\t[default: %s]' % str(default) 138 | message += '{:>18}: {:<15}{}\n'.format(str(key), str(value), comment) 139 | message += '----------------------------------------------------------------\n' 140 | print(message) 141 | 142 | # Save the running parameters setting in the disk 143 | experiment_dir = os.path.join(param.checkpoints_dir, param.experiment_name) 144 | util.mkdir(experiment_dir) 145 | file_name = os.path.join(experiment_dir, 'cmd_parameters.txt') 146 | with open(file_name, 'w') as param_file: 147 | now = time.strftime('%c') 148 | param_file.write('{:s}\n'.format(now)) 149 | param_file.write(message) 150 | param_file.write('\n') 151 | 152 | def parse(self): 153 | """ 154 | Parse the parameters of our project. Set up GPU device. Print the welcome words and list parameters in the console. 155 | """ 156 | param = self.get_params() # get the parameters to the object param 157 | 158 | param.isTrain = self.isTrain # store this if for training or testing 159 | param.isTest = self.isTest 160 | 161 | # Print welcome words and command line parameters 162 | self.print_params(param) 163 | 164 | # Set the internal parameters 165 | # add_channel: add one extra dimension of channel for the input data, used for convolution layer 166 | # if param.netG == 'unet_1d' or param.netG == 'unet_s_1d' or param.netD == 'multi_conv_1d': 167 | # param.add_channel = True 168 | # else: 169 | # param.add_channel = False 170 | param.add_channel = True 171 | # ch_separate: separate the DNA methylation matrix base on the chromosome 172 | if param.netG == 'fc_ved_sep' or param.netG == 'fcg_sep': 173 | param.ch_separate = True 174 | else: 175 | param.ch_separate = False 176 | 177 | # Set up GPU 178 | str_gpu_ids = param.gpu_ids.split(',') 179 | param.gpu_ids = [] 180 | for str_gpu_id in str_gpu_ids: 181 | int_gpu_id = int(str_gpu_id) 182 | if int_gpu_id >= 0: 183 | param.gpu_ids.append(int_gpu_id) 184 | if len(param.gpu_ids) > 0: 185 | torch.cuda.set_device(param.gpu_ids[0]) 186 | 187 | self.param = param 188 | return self.param 189 | -------------------------------------------------------------------------------- /params/test_params.py: -------------------------------------------------------------------------------- 1 | from .basic_params import BasicParams 2 | 3 | 4 | class TestParams(BasicParams): 5 | """ 6 | This class is a son class of BasicParams. 7 | This class includes parameters for testing and parameters inherited from the father class. 8 | """ 9 | def initialize(self, parser): 10 | parser = BasicParams.initialize(self, parser) 11 | 12 | # Testing parameters 13 | parser.add_argument('--start_epoch', type=int, default=1, 14 | help='start epoch number for testing') 15 | parser.add_argument('--end_epoch', type=int, default=1, 16 | help='end epoch number for testing') 17 | 18 | # Logging and visualization 19 | parser.add_argument('--print_freq', type=int, default=10, 20 | help='frequency of showing testing results on console') 21 | parser.add_argument('--save_fake', action='store_true', 22 | help='save the fake omics data to disc') 23 | 24 | self.isTrain = False 25 | self.isTest = True 26 | return parser 27 | -------------------------------------------------------------------------------- /params/train_params.py: -------------------------------------------------------------------------------- 1 | from .basic_params import BasicParams 2 | 3 | 4 | class TrainParams(BasicParams): 5 | """ 6 | This class is a son class of BasicParams. 7 | This class includes parameters for training and parameters inherited from the father class. 8 | """ 9 | def initialize(self, parser): 10 | parser = BasicParams.initialize(self, parser) 11 | 12 | # Training parameters 13 | parser.add_argument('--GAN_mode', type=str, default='vanilla', 14 | help='The type of GAN objective. Vanilla GAN loss is the cross-entropy objective used in the original GAN paper. [vanilla | lsgan]') 15 | parser.add_argument('--lr_G', type=float, default=0.0002, 16 | help='initial learning rate for the generator') 17 | parser.add_argument('--lr_D', type=float, default=0.0002, 18 | help='initial learning rate for the discriminator') 19 | parser.add_argument('--beta1', type=float, default=0.5, 20 | help='momentum term of adam') 21 | parser.add_argument('--lr_policy', type=str, default='linear', 22 | help='The learning rate policy for the scheduler. [linear | step | plateau | cosine]') 23 | parser.add_argument('--epoch_count', type=int, default=1, 24 | help='the starting epoch count, default start from 1') 25 | parser.add_argument('--epoch_num', type=int, default=100, 26 | help='Number of epoch using starting learning rate') 27 | parser.add_argument('--epoch_num_decay', type=int, default=0, 28 | help='Number of epoch to linearly decay learning rate to zero') 29 | parser.add_argument('--decay_step_size', type=int, default=50, 30 | help='The original learning rate multiply by a gamma every decay_step_size epoch') 31 | 32 | # Network saving and loading parameters 33 | parser.add_argument('--continue_train', action='store_true', 34 | help='load the latest model and continue training ') 35 | parser.add_argument('--save_model', action='store_true', 36 | help='save the model during training') 37 | parser.add_argument('--save_epoch_freq', type=int, default=-1, 38 | help='frequency of saving checkpoints at the end of epochs, -1 means only save the last epoch') 39 | 40 | # Logging and visualization 41 | parser.add_argument('--print_freq', type=int, default=1, 42 | help='frequency of showing training results on console') 43 | 44 | self.isTrain = True 45 | self.isTest = False 46 | return parser 47 | -------------------------------------------------------------------------------- /params/train_test_param.py: -------------------------------------------------------------------------------- 1 | from .basic_params import BasicParams 2 | 3 | 4 | class TrainTestParams(BasicParams): 5 | """ 6 | This class is a son class of BasicParams. 7 | This class includes parameters for both training and testing. 8 | """ 9 | def initialize(self, parser): 10 | parser = BasicParams.initialize(self, parser) 11 | 12 | # Training parameters 13 | parser.add_argument('--GAN_mode', type=str, default='vanilla', 14 | help='The type of GAN objective. Vanilla GAN loss is the cross-entropy objective used in the original GAN paper. [vanilla | lsgan]') 15 | parser.add_argument('--lr_G', type=float, default=0.0002, 16 | help='initial learning rate for the generator') 17 | parser.add_argument('--lr_D', type=float, default=0.0002, 18 | help='initial learning rate for the discriminator') 19 | parser.add_argument('--beta1', type=float, default=0.5, 20 | help='momentum term of adam') 21 | parser.add_argument('--lr_policy', type=str, default='linear', 22 | help='The learning rate policy for the scheduler. [linear | step | plateau | cosine]') 23 | parser.add_argument('--epoch_count', type=int, default=1, 24 | help='the starting epoch count, default start from 1') 25 | parser.add_argument('--epoch_num', type=int, default=100, 26 | help='Number of epoch using starting learning rate') 27 | parser.add_argument('--epoch_num_decay', type=int, default=0, 28 | help='Number of epoch to linearly decay learning rate to zero') 29 | parser.add_argument('--decay_step_size', type=int, default=50, 30 | help='The original learning rate multiply by a gamma every decay_step_size epoch') 31 | 32 | # Network saving and loading parameters 33 | parser.add_argument('--continue_train', action='store_true', 34 | help='load the latest model and continue training') 35 | parser.add_argument('--save_model', action='store_true', 36 | help='save the model during training') 37 | parser.add_argument('--save_epoch_freq', type=int, default=-1, 38 | help='frequency of saving checkpoints at the end of epochs, -1 means only save the last epoch') 39 | 40 | # Logging and visualization 41 | parser.add_argument('--print_freq', type=int, default=1, 42 | help='frequency of showing results on console') 43 | parser.add_argument('--save_fake', action='store_true', 44 | help='save the fake omics data to disc') 45 | 46 | # Dataset parameters 47 | parser.add_argument('--stratify', action='store_true', 48 | help='apply the stratified mode in train/test split if set true') 49 | parser.add_argument('--train_ratio', type=float, default=0.8, 50 | help='ratio of training set in the full dataset') 51 | parser.add_argument('--test_ratio', type=float, default=0.2, 52 | help='ratio of testing set in the full dataset') 53 | 54 | self.isTrain = True 55 | self.isTest = True 56 | return parser 57 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch>=1.2.0 2 | tensorboard>=1.10.0 3 | scikit-learn>=1.0.2 4 | pandas>=1.3.5 5 | prefetch_generator>=1.0.0 6 | tables>=3.6.0 7 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | """ 2 | The testing part for OmiTrans 3 | """ 4 | import time 5 | from util import util 6 | from params.test_params import TestParams 7 | from datasets import create_single_dataloader 8 | from models import create_model 9 | from util.visualizer import Visualizer 10 | 11 | 12 | if __name__ == '__main__': 13 | # Get testing parameter 14 | param = TestParams().parse() 15 | if param.deterministic: 16 | util.setup_seed(param.seed) 17 | 18 | # Dataset related 19 | dataloader, param.sample_list = create_single_dataloader(param, shuffle=False) # No shuffle for testing 20 | print('The size of testing set is {}'.format(len(dataloader))) 21 | # Get sample list of the dataset 22 | sample_list = dataloader.get_sample_list() 23 | # Get feature list of the dataset 24 | feature_list_A = dataloader.get_feature_list_A() 25 | # Get the dimension of two omics types 26 | param.A_dim = dataloader.get_A_dim() 27 | param.B_dim = dataloader.get_B_dim() 28 | print('The dimension of omics type A is %d' % param.A_dim) 29 | print('The dimension of omics type B is %d' % param.B_dim) 30 | if param.zo_norm: 31 | param.target_min = dataloader.get_values_min() 32 | param.target_max = dataloader.get_values_max() 33 | 34 | # Model related 35 | model = create_model(param) # Create a model given param.model and other parameters 36 | model.set_eval() 37 | visualizer = Visualizer(param) # Create a visualizer to print results 38 | 39 | # TESTING 40 | model.setup(param) # load saved networks for testing 41 | metrics_acc = model.init_metrics_dict() # Initialize the metrics dictionary 42 | if param.save_fake: # Initialize the fake array during the last epoch 43 | fake_dict = model.init_fake_dict() 44 | test_start_time = time.time() # Start time of testing 45 | 46 | for i, data in enumerate(dataloader): 47 | dataset_size = len(dataloader) 48 | model.set_input(data) # Unpack input data from the output dictionary of the dataloader 49 | model.test() # Run forward to get the fake omics data 50 | model.update_metrics_dict(metrics_acc) # Update the metrics dictionary 51 | if param.save_fake: # Update the fake array during the last epoch 52 | fake_dict = model.update_fake_dict(fake_dict) 53 | visualizer.print_test_log(param.epoch_to_load, i, metrics_acc, param.batch_size, dataset_size) 54 | test_time = time.time() - test_start_time 55 | # Save average metrics of this epoch to the disk 56 | visualizer.print_test_summary(param.epoch_to_load, metrics_acc, test_time) 57 | 58 | if param.save_fake: # Save the fake omics data 59 | visualizer.save_fake_omics(fake_dict, sample_list, feature_list_A) 60 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | """ 2 | The training part for OmiTrans 3 | """ 4 | import time 5 | from util import util 6 | from params.train_params import TrainParams 7 | from datasets import create_single_dataloader 8 | from models import create_model 9 | from util.visualizer import Visualizer 10 | 11 | 12 | if __name__ == '__main__': 13 | full_start_time = time.time() 14 | 15 | # Get parameters 16 | param = TrainParams().parse() 17 | if param.deterministic: 18 | util.setup_seed(param.seed) 19 | 20 | # Dataset related 21 | dataloader, sample_list = create_single_dataloader(param, shuffle=True, enable_drop_last=True) 22 | print('The size of training set is {}'.format(len(dataloader))) 23 | # Get the dimension of two omics types 24 | param.A_dim = dataloader.get_A_dim() 25 | param.B_dim = dataloader.get_B_dim() 26 | print('The dimension of omics type A is %d' % param.A_dim) 27 | print('The dimension of omics type B is %d' % param.B_dim) 28 | if param.zo_norm: 29 | param.target_min = dataloader.get_values_min() 30 | param.target_max = dataloader.get_values_max() 31 | 32 | # Model related 33 | model = create_model(param) # Create a model given param.model and other parameters 34 | model.setup(param) # Regular setup for the model: load and print networks, create schedulers 35 | visualizer = Visualizer(param) # Create a visualizer to print results 36 | 37 | # Start the epoch loop 38 | for epoch in range(param.epoch_count, param.epoch_num + param.epoch_num_decay + 1): # outer loop for different epochs 39 | epoch_start_time = time.time() # Start time of this epoch 40 | 41 | # Training 42 | model.set_train() # Set train mode for training 43 | iter_load_start_time = time.time() # Start time of data loading for this iteration 44 | losses_acc = model.init_losses_dict() # Initialize the losses dictionary 45 | metrics_acc_train = model.init_metrics_dict() # Initialize the metrics dictionary 46 | 47 | # Start training loop 48 | for i, data in enumerate(dataloader): # Inner loop for different iteration within one epoch 49 | dataset_size = len(dataloader) 50 | actual_batch_size = len(data['index']) 51 | iter_start_time = time.time() # Timer for computation per iteration 52 | if i % param.print_freq == 0: 53 | load_time = iter_start_time - iter_load_start_time # Data loading time for this iteration 54 | model.set_input(data) # Unpack input data from the output dictionary of the dataloader 55 | model.update() # Calculate losses, gradients and update network parameters 56 | model.update_losses_dict(losses_acc) # Update the losses dictionary 57 | model.update_metrics_dict(metrics_acc_train) # Update the metrics dictionary 58 | if i % param.print_freq == 0: # Print training losses and save logging information to the disk 59 | comp_time = time.time() - iter_start_time # Computational time for this iteration 60 | visualizer.print_train_log(epoch, i, losses_acc, metrics_acc_train, load_time / actual_batch_size, comp_time / actual_batch_size, param.batch_size, dataset_size) 61 | iter_load_start_time = time.time() 62 | 63 | if param.save_model: 64 | if param.save_epoch_freq == -1: # Only save networks during last epoch 65 | if epoch == param.epoch_num + param.epoch_num_decay: 66 | print('Saving the model at the end of epoch {:d}'.format(epoch)) 67 | model.save_networks(str(epoch)) 68 | elif epoch % param.save_epoch_freq == 0: # Save both the generator and the discriminator every epochs 69 | print('Saving the model at the end of epoch {:d}'.format(epoch)) 70 | # model.save_networks('latest') 71 | model.save_networks(str(epoch)) 72 | 73 | # total num of epoch for the training phase 74 | total_epoch = param.epoch_num + param.epoch_num_decay 75 | # Running time for this epoch 76 | train_time = time.time() - epoch_start_time 77 | visualizer.print_train_summary(epoch, losses_acc, metrics_acc_train, train_time) 78 | model.update_learning_rate() # update learning rates at the end of each epoch 79 | 80 | full_time = time.time() - full_start_time 81 | print('Full running time: {:.3f}s'.format(full_time)) 82 | -------------------------------------------------------------------------------- /train_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Training and testing for OmiTrans 3 | """ 4 | import time 5 | from util import util 6 | from params.train_test_param import TrainTestParams 7 | from datasets import create_separate_dataloader 8 | from models import create_model 9 | from util.visualizer import Visualizer 10 | 11 | 12 | if __name__ == '__main__': 13 | full_start_time = time.time() 14 | 15 | # Get parameters 16 | param = TrainTestParams().parse() 17 | if param.deterministic: 18 | util.setup_seed(param.seed) 19 | 20 | # Dataset related 21 | full_dataloader, train_dataloader, val_dataloader, test_dataloader = create_separate_dataloader(param) 22 | dataset_size_train = len(train_dataloader) # Get the number of data points in the training set. 23 | print('The number of training omics data points is %d' % dataset_size_train) 24 | # Get sample list of the dataset 25 | sample_list = full_dataloader.get_sample_list() 26 | # Get feature list of the dataset 27 | feature_list_A = full_dataloader.get_feature_list_A() 28 | # Get the dimension of two omics types 29 | param.A_dim = full_dataloader.get_A_dim() 30 | param.B_dim = full_dataloader.get_B_dim() 31 | print('The dimension of omics type A is %d' % param.A_dim) 32 | print('The dimension of omics type B is %d' % param.B_dim) 33 | if param.zo_norm: 34 | param.target_min = full_dataloader.get_values_min() 35 | param.target_max = full_dataloader.get_values_max() 36 | 37 | # Model related 38 | model = create_model(param) # Create a model given param.model and other parameters 39 | model.setup(param) # Regular setup for the model: load and print networks, create schedulers 40 | visualizer = Visualizer(param) # Create a visualizer to print results 41 | 42 | # Start the epoch loop 43 | for epoch in range(param.epoch_count, param.epoch_num + param.epoch_num_decay + 1): # outer loop for different epochs 44 | epoch_start_time = time.time() # Start time of this epoch 45 | 46 | # Training 47 | model.set_train() # Set train mode for training 48 | iter_load_start_time = time.time() # Start time of data loading for this iteration 49 | losses_acc = model.init_losses_dict() # Initialize the losses dictionary 50 | metrics_acc_train = model.init_metrics_dict() # Initialize the metrics dictionary 51 | 52 | # Start training loop 53 | for i, data in enumerate(train_dataloader): # Inner loop for different iteration within one epoch 54 | dataset_size = len(train_dataloader) 55 | actual_batch_size = len(data['index']) 56 | iter_start_time = time.time() # Timer for computation per iteration 57 | if i % param.print_freq == 0: 58 | load_time = iter_start_time - iter_load_start_time # Data loading time for this iteration 59 | model.set_input(data) # Unpack input data from the output dictionary of the dataloader 60 | model.update() # Calculate losses, gradients and update network parameters 61 | model.update_losses_dict(losses_acc) # Update the losses dictionary 62 | model.update_metrics_dict(metrics_acc_train) # Update the metrics dictionary 63 | if i % param.print_freq == 0: # Print training losses and save logging information to the disk 64 | comp_time = time.time() - iter_start_time # Computational time for this iteration 65 | visualizer.print_train_log(epoch, i, losses_acc, metrics_acc_train, load_time / actual_batch_size, comp_time / actual_batch_size, param.batch_size, dataset_size) 66 | iter_load_start_time = time.time() 67 | 68 | if param.save_model: 69 | if param.save_epoch_freq == -1: # Only save networks during last epoch 70 | if epoch == param.epoch_num + param.epoch_num_decay: 71 | print('Saving the model at the end of epoch {:d}'.format(epoch)) 72 | model.save_networks(str(epoch)) 73 | elif epoch % param.save_epoch_freq == 0: # Save both the generator and the discriminator every epochs 74 | print('Saving the model at the end of epoch {:d}'.format(epoch)) 75 | # model.save_networks('latest') 76 | model.save_networks(str(epoch)) 77 | 78 | # total num of epoch for the training phase 79 | total_epoch = param.epoch_num + param.epoch_num_decay 80 | # Running time for this epoch 81 | train_time = time.time() - epoch_start_time 82 | visualizer.print_train_summary(epoch, losses_acc, metrics_acc_train, train_time) 83 | model.update_learning_rate() # update learning rates at the end of each epoch 84 | 85 | # Testing 86 | model.set_eval() # Set eval mode for testing 87 | metrics_acc_test = model.init_metrics_dict() # Initialize the metrics dictionary 88 | if epoch == param.epoch_num + param.epoch_num_decay and param.save_fake: # Initialize the fake array during the last epoch 89 | fake_dict = model.init_fake_dict() 90 | test_start_time = time.time() # Start time of testing 91 | 92 | # Start testing loop 93 | for i, data in enumerate(test_dataloader): 94 | dataset_size = len(test_dataloader) 95 | model.set_input(data) # Unpack input data from the output dictionary of the dataloader 96 | model.test() # Run forward to get the fake omics data 97 | model.update_metrics_dict(metrics_acc_test) # Update the metrics dictionary 98 | if epoch == param.epoch_num + param.epoch_num_decay and param.save_fake: # Update the fake array during the last epoch 99 | fake_dict = model.update_fake_dict(fake_dict) 100 | if i % param.print_freq == 0: # Print testing losses 101 | visualizer.print_test_log(epoch, i, metrics_acc_test, param.batch_size, dataset_size) 102 | test_time = time.time() - test_start_time 103 | # Save average metrics of this epoch to the disk 104 | visualizer.print_test_summary(epoch, metrics_acc_test, test_time) 105 | if epoch == param.epoch_num + param.epoch_num_decay and param.save_fake: # Save the fake omics data 106 | visualizer.save_fake_omics(fake_dict, sample_list, feature_list_A) 107 | 108 | full_time = time.time() - full_start_time 109 | print('Full running time: {:.3f}s'.format(full_time)) 110 | -------------------------------------------------------------------------------- /util/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | This package includes some useful helper functions. 3 | """ 4 | -------------------------------------------------------------------------------- /util/preprocess.py: -------------------------------------------------------------------------------- 1 | """ 2 | Contain some omics data preprocess functions 3 | """ 4 | import pandas as pd 5 | 6 | 7 | def separate_B(B_df_single): 8 | """ 9 | Separate the DNA methylation dataframe into subsets according to their targeting chromosomes 10 | 11 | Parameters: 12 | B_df_single(DataFrame) -- a dataframe that contains the single DNA methylation matrix 13 | 14 | Return: 15 | B_df_list(list) -- a list with 23 subset dataframe 16 | B_dim(list) -- the dims of each chromosome 17 | """ 18 | anno = pd.read_csv('./anno/B_anno.csv', dtype={'CHR': str}, index_col=0) 19 | anno_contain = anno.loc[B_df_single.index, :] 20 | print('Separating B.tsv according the targeting chromosome...') 21 | B_df_list, B_dim_list = [], [] 22 | ch_id = list(range(1, 23)) 23 | ch_id.append('X') 24 | for ch in ch_id: 25 | ch_index = anno_contain[anno_contain.CHR == str(ch)].index 26 | ch_df = B_df_single.loc[ch_index, :] 27 | B_df_list.append(ch_df) 28 | B_dim_list.append(len(ch_df)) 29 | 30 | return B_df_list, B_dim_list 31 | -------------------------------------------------------------------------------- /util/util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Contain some simple helper functions 3 | """ 4 | import os 5 | import shutil 6 | import torch 7 | import random 8 | import numpy as np 9 | 10 | 11 | def mkdir(path): 12 | """ 13 | Create a empty directory in the disk if it didn't exist 14 | 15 | Parameters: 16 | path(str) -- a directory path we would like to create 17 | """ 18 | if not os.path.exists(path): 19 | os.mkdir(path) 20 | 21 | 22 | def clear_dir(path): 23 | """ 24 | delete all files in a path 25 | 26 | Parameters: 27 | path(str) -- a directory path that we would like to delete all files in it 28 | """ 29 | if os.path.exists(path): 30 | shutil.rmtree(path) 31 | os.mkdir(path) 32 | 33 | 34 | def setup_seed(seed): 35 | """ 36 | setup seed to make the experiments deterministic 37 | 38 | Parameters: 39 | seed(int) -- the random seed 40 | """ 41 | torch.manual_seed(seed) 42 | torch.cuda.manual_seed_all(seed) 43 | np.random.seed(seed) 44 | random.seed(seed) 45 | torch.backends.cudnn.deterministic = True 46 | -------------------------------------------------------------------------------- /util/visualizer.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import numpy as np 4 | import pandas as pd 5 | from util import util 6 | from datasets import save_file 7 | from torch.utils.tensorboard import SummaryWriter 8 | 9 | 10 | class Visualizer: 11 | """ 12 | This class print/save logging information 13 | """ 14 | 15 | def __init__(self, param): 16 | """ 17 | Initialize the Visualizer class 18 | """ 19 | self.param = param 20 | tb_dir = os.path.join(param.checkpoints_dir, param.experiment_name, 'tb_log') 21 | util.mkdir(tb_dir) 22 | 23 | if param.isTrain: 24 | # Create a logging file to store training losses 25 | self.train_log_filename = os.path.join(param.checkpoints_dir, param.experiment_name, 'train_log.txt') 26 | with open(self.train_log_filename, 'a') as log_file: 27 | now = time.strftime('%c') 28 | log_file.write('----------------------- Training Log ({:s}) -----------------------\n'.format(now)) 29 | 30 | self.train_summary_filename = os.path.join(param.checkpoints_dir, param.experiment_name, 'train_summary.txt') 31 | with open(self.train_summary_filename, 'a') as log_file: 32 | now = time.strftime('%c') 33 | log_file.write('----------------------- Training Summary ({:s}) -----------------------\n'.format(now)) 34 | 35 | # Create log folder for TensorBoard 36 | tb_train_dir = os.path.join(param.checkpoints_dir, param.experiment_name, 'tb_log', 'train') 37 | util.mkdir(tb_train_dir) 38 | util.clear_dir(tb_train_dir) 39 | 40 | # Create TensorBoard writer 41 | self.train_writer = SummaryWriter(log_dir=tb_train_dir) 42 | 43 | if param.isTest: 44 | # Create a logging file to store testing metrics 45 | self.test_log_filename = os.path.join(param.checkpoints_dir, param.experiment_name, 'test_log.txt') 46 | with open(self.test_log_filename, 'a') as log_file: 47 | now = time.strftime('%c') 48 | log_file.write('----------------------- Testing Log ({:s}) -----------------------\n'.format(now)) 49 | 50 | self.test_summary_filename = os.path.join(param.checkpoints_dir, param.experiment_name, 'test_summary.txt') 51 | with open(self.test_summary_filename, 'a') as log_file: 52 | now = time.strftime('%c') 53 | log_file.write('----------------------- Testing Summary ({:s}) -----------------------\n'.format(now)) 54 | 55 | # Create log folder for TensorBoard 56 | tb_test_dir = os.path.join(param.checkpoints_dir, param.experiment_name, 'tb_log', 'test') 57 | util.mkdir(tb_test_dir) 58 | util.clear_dir(tb_test_dir) 59 | 60 | # Create TensorBoard writer 61 | self.test_writer = SummaryWriter(log_dir=tb_test_dir) 62 | 63 | def print_train_log(self, epoch, iteration, losses, metrics, load_time, comp_time, batch_size, dataset_size): 64 | """ 65 | print train log on console and save the message to the disk 66 | 67 | Parameters: 68 | epoch (int) -- current epoch 69 | iteration (int) -- current training iteration during this epoch 70 | losses (OrderedDict) -- training losses stored in the ordered dict 71 | metrics (OrderedDict) -- metrics stored in the ordered dict 72 | load_time (float) -- data loading time per data point (normalized by batch_size) 73 | comp_time (float) -- computational time per data point (normalized by batch_size) 74 | batch_size (int) -- batch size of training 75 | dataset_size (int) -- size of the training dataset 76 | """ 77 | data_point_covered = min((iteration + 1) * batch_size, dataset_size) 78 | message = '[TRAIN] [Epoch: {:3d} Iter: {:4d} Load_t: {:.3f} Comp_t: {:.3f}] '.format(epoch, data_point_covered, load_time, comp_time) 79 | for name, loss in losses.items(): 80 | message += '{:s}: {:.4f} '.format(name, loss[-1]) 81 | for name, metric in metrics.items(): 82 | message += '{:s}: {:.4f} '.format(name, metric[-1]) 83 | 84 | print(message) # print the message 85 | 86 | with open(self.train_log_filename, 'a') as log_file: 87 | log_file.write(message + '\n') # save the message 88 | 89 | def print_train_summary(self, epoch, losses, metrics, train_time): 90 | """ 91 | print the summary of this training epoch 92 | 93 | Parameters: 94 | epoch (int) -- epoch number of this training model 95 | losses (OrderedDict) -- the losses dictionary 96 | metrics (OrderedDict) -- the metrics dictionary 97 | train_time (float) -- time used for training this epoch 98 | """ 99 | write_message = '{:s}\t'.format(str(epoch)) 100 | print_message = '[TRAIN] [Epoch: {:4d}] '.format(int(epoch)) 101 | 102 | for name, loss in losses.items(): 103 | write_message += '{:.6f}\t'.format(np.mean(loss)) 104 | print_message += name + ': {:.6f} '.format(np.mean(loss)) 105 | self.train_writer.add_scalar(name, np.mean(loss), epoch) 106 | for name, metric in metrics.items(): 107 | write_message += '{:.6f}\t'.format(np.mean(metric)) 108 | print_message += name + ': {:.6f} '.format(np.mean(metric)) 109 | self.train_writer.add_scalar(name, np.mean(metric), epoch) 110 | 111 | with open(self.train_summary_filename, 'a') as log_file: 112 | log_file.write(write_message + '\n') 113 | 114 | train_time_msg = 'Training time used: {:.3f}s'.format(train_time) 115 | print_message += '\n' + train_time_msg 116 | print(print_message) 117 | with open(self.train_log_filename, 'a') as log_file: 118 | log_file.write(train_time_msg + '\n') 119 | 120 | def print_test_log(self, epoch, iteration, metrics, batch_size, dataset_size): 121 | """ 122 | print performance metrics of this iteration on console and save the message to the disk 123 | 124 | Parameters: 125 | epoch (int) -- epoch number of this testing model 126 | iteration (int) -- current testing iteration during this epoch 127 | metrics (OrderedDict) -- testing metrics stored in the dictionary 128 | batch_size (int) -- batch size of testing 129 | dataset_size (int) -- size of the testing dataset 130 | """ 131 | data_point_covered = min((iteration + 1) * batch_size, dataset_size) 132 | message = '[TEST] [Epoch: {:3d} Iter: {:4d}] '.format(int(epoch), data_point_covered) 133 | for name, metric in metrics.items(): 134 | message += '{:s}: {:.4f} '.format(name, metric[-1]) 135 | 136 | print(message) 137 | 138 | with open(self.test_log_filename, 'a') as log_file: 139 | log_file.write(message + '\n') 140 | 141 | def print_test_summary(self, epoch, metrics, test_time): 142 | """ 143 | print the summary of this testing epoch 144 | 145 | Parameters: 146 | epoch (int) -- epoch number of this testing model 147 | metrics (OrderedDict) -- the metrics dictionary 148 | test_time (float) -- time used for testing this epoch 149 | """ 150 | write_message = '{:s}\t'.format(str(epoch)) 151 | print_message = '[TEST] [Epoch: {:4d}] '.format(int(epoch)) 152 | 153 | for name, metric in metrics.items(): 154 | write_message += '{:.6f}\t'.format(np.mean(metric)) 155 | print_message += name + ': {:.6f} '.format(np.mean(metric)) 156 | self.test_writer.add_scalar(name, np.mean(metric), epoch) 157 | 158 | with open(self.test_summary_filename, 'a') as log_file: 159 | log_file.write(write_message + '\n') 160 | 161 | test_time_msg = 'Testing time used: {:.3f}s'.format(test_time) 162 | print_message += '\n' + test_time_msg 163 | print(print_message) 164 | with open(self.test_log_filename, 'a') as log_file: 165 | log_file.write(test_time_msg + '\n') 166 | 167 | def save_fake_omics(self, fake_dict, sample_list, feature_list): 168 | """ 169 | save the fake omics data to disc 170 | 171 | Parameters: 172 | fake_dict (OrderedDict)) -- the fake omics data and the corresponding index 173 | sample_list (ndarray) -- the sample list for the input data 174 | feature_list (ndarray) -- the feature list of the generated omics data 175 | """ 176 | output_sample_list = sample_list[fake_dict['index'].astype(int)] 177 | fake_df = pd.DataFrame(data=fake_dict['fake'].T, index=feature_list, columns=output_sample_list) 178 | print('Saving generated omics file...') 179 | save_file(self.param, fake_df, 'fake_A') 180 | --------------------------------------------------------------------------------