├── .gitignore ├── LICENSE ├── README-zh.md ├── README.md ├── data_function.py ├── hparam.py ├── loss_function.py ├── main.py ├── models ├── three_d │ ├── densenet3d.py │ ├── resnet3d.py │ └── resnext3d.py └── two_d │ ├── alexnet.py │ ├── densenet.py │ ├── googlenet.py │ ├── mobilenet.py │ ├── nasnet.py │ ├── resnet.py │ ├── resnext.py │ └── vggnet.py └── utils ├── convolution.py ├── dilation.py ├── metric.py └── residual.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Kangneng Zhou 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-zh.md: -------------------------------------------------------------------------------- 1 | # Pytorch Medical Segmentation 2 | 英文版请戳：这里！
3 | 4 | 5 | ## 最近的更新 6 | * 2021.7.7 训练和测试代码已经发布 7 | 8 | 9 | ## 环境要求 10 | * pytorch1.7 11 | * torchio<=0.18.20 12 | * python>=3.6 13 | 14 | ## 通知 15 | * 您可以修改**hparam.py**文件来确定是2D分类还是3D分类以及是否可以进行多分类。 16 | * 我们几乎提供了所有的2D和3D分类的算法。 17 | * 本项目兼容几乎所有的医学数据格式(例如 png, nii.gz, nii, mhd, nrrd, ...)，修改**hparam.py**的**fold_arch**即可。 18 | * 如果您想进行**多分类**分割，请自行修改**data_function.py**的对应代码。我不能确定您的具体分类数。 19 | 20 | 21 | 22 | ## 准备您的数据 23 | ### 例 24 | 如果您的source文件夹如下排列 : 25 | ``` 26 | categpry-0 27 | ├── source_1.png 28 | ├── source_2.png 29 | ├── source_3.png 30 | └── ... 31 | ``` 32 | 33 | ``` 34 | categpry-1 35 | ├── source_1.png 36 | ├── source_2.png 37 | ├── source_3.png 38 | └── ... 39 | ``` 40 | 41 | 您应该修改 **fold_arch** 为 **\*/\*.png**, **source_train_0_dir** 为 **categpry-0** 并修改 **source_train_1_dir** 为 **categpry-1** in **hparam.py** 42 | 43 | 44 | 45 | ## 训练 46 | * 不使用预训练模型 47 | ``` 48 | set hparam.train_or_test to 'train' 49 | python main.py 50 | ``` 51 | * 使用预训练模型 52 | ``` 53 | set hparam.train_or_test to 'train' 54 | set hparam.ckpt to True 55 | python main.py 56 | ``` 57 | 58 | ## Inference 59 | * 测试 60 | ``` 61 | set hparam.train_or_test to 'test' 62 | python main.py 63 | ``` 64 | 65 | 66 | 67 | ## Done 68 | ### Network 69 | * 2D 70 | - [x] alexnet 71 | - [x] densenet 72 | - [x] googlenet 73 | - [x] mobilenet 74 | - [x] nasnet 75 | - [x] resnet 76 | - [x] resnext 77 | - [x] vggnet 78 | * 3D 79 | - [x] densenet3d 80 | - [x] resnet3d 81 | - [x] resnext3d 82 | 83 | ## TODO 84 | - [ ] dataset 85 | - [ ] benchmark 86 | 87 | ## By The Way 88 | 这个项目并不完美，还存在很多问题。如果您正在使用这个项目，并想给作者一些反馈，您可以给[Kangneng Zhou](elliszkn@163.com)发邮件，或添加他的**微信**：ellisgege666 89 | 90 | ## 致谢 91 | 这个项目是一个非官方PyTorch实现的3D和2D医学分类，高度依赖于[pytorch-cifar100](https://github.com/weiaicunzai/pytorch-cifar100)和[torchio](https://github.com/fepegar/torchio)。感谢上述项目。感谢[Cheng Chen](b20170310@xs.ustb.edu.cn)和[Weili Jiang](1379252229@qq.com)对我的帮助。 92 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pytorch Medical Classification 2 | Read Chinese Introduction：Here！
3 | 4 | ## Recent Updates 5 | * 2021.7.7 The train and test codes are released. 6 | 7 | 8 | ## Requirements 9 | * pytorch1.7 10 | * torchio<=0.18.20 11 | * python>=3.6 12 | 13 | ## Notice 14 | * You can modify **hparam.py** to determine whether 2D or 3D classification and whether multicategorization is possible. 15 | * We provide algorithms for almost all 2D and 3D classification. 16 | * This repository is compatible with almost all medical data formats(e.g. png, nii.gz, nii, mhd, nrrd, ...), by modifying **fold_arch** in **hparam.py** of the config. 17 | * If you want to use a **multi-category** program, please modify the corresponding codes in **data_function.py** by yourself. I cannot identify your specific categories. 18 | 19 | ## Prepare Your Dataset 20 | ### Example 21 | if your source dataset is : 22 | ``` 23 | categpry-0 24 | ├── source_1.png 25 | ├── source_2.png 26 | ├── source_3.png 27 | └── ... 28 | ``` 29 | 30 | ``` 31 | categpry-1 32 | ├── source_1.png 33 | ├── source_2.png 34 | ├── source_3.png 35 | └── ... 36 | ``` 37 | 38 | 39 | then your should modify **fold_arch** as **\*.png**, **source_train_0_dir** as **categpry-0** and **source_train_1_dir** as **categpry-1** in **hparam.py** 40 | 41 | 42 | 43 | ## Training 44 | * without pretrained-model 45 | ``` 46 | set hparam.train_or_test to 'train' 47 | python main.py 48 | ``` 49 | * with pretrained-model 50 | ``` 51 | set hparam.train_or_test to 'train' 52 | set hparam.ckpt to True 53 | python main.py 54 | ``` 55 | 56 | ## Inference 57 | * testing 58 | ``` 59 | set hparam.train_or_test to 'test' 60 | python main.py 61 | ``` 62 | 63 | 64 | ## Done 65 | ### Network 66 | * 2D 67 | - [x] alexnet 68 | - [x] densenet 69 | - [x] googlenet 70 | - [x] mobilenet 71 | - [x] nasnet 72 | - [x] resnet 73 | - [x] resnext 74 | - [x] vggnet 75 | * 3D 76 | - [x] densenet3d 77 | - [x] resnet3d 78 | - [x] resnext3d 79 | 80 | 81 | 82 | ## TODO 83 | - [ ] dataset 84 | - [ ] benchmark 85 | 86 | 87 | ## By The Way 88 | This project is not perfect and there are still many problems. If you are using this project and would like to give the author some feedbacks, you can send [Kangneng Zhou](elliszkn@163.com) an email, his **wechat** number is: ellisgege666 89 | 90 | ## Acknowledgements 91 | This repository is an unoffical PyTorch implementation of Medical segmentation in 3D and 2D and highly based on [pytorch-cifar100](https://github.com/weiaicunzai/pytorch-cifar100) and [torchio](https://github.com/fepegar/torchio).Thank you for the above repo. Thank you to [Cheng Chen](b20170310@xs.ustb.edu.cn) and [Weili Jiang](1379252229@qq.com) for all the help I received. 92 | -------------------------------------------------------------------------------- /data_function.py: -------------------------------------------------------------------------------- 1 | from glob import glob 2 | from os.path import dirname, join, basename, isfile 3 | import sys 4 | sys.path.append('./') 5 | import csv 6 | import torch 7 | from medpy.io import load 8 | import numpy as np 9 | from PIL import Image 10 | from torch import nn 11 | import torch.nn.functional as F 12 | import random 13 | import torchio as tio 14 | from torchio import AFFINE, DATA 15 | import torchio 16 | from torchio import ScalarImage, LabelMap, Subject, SubjectsDataset, Queue 17 | from torchio.data import UniformSampler 18 | from torchio.transforms import ( 19 | RandomFlip, 20 | RandomAffine, 21 | RandomElasticDeformation, 22 | RandomNoise, 23 | RandomMotion, 24 | RandomBiasField, 25 | RescaleIntensity, 26 | Resample, 27 | ToCanonical, 28 | ZNormalization, 29 | CropOrPad, 30 | HistogramStandardization, 31 | OneOf, 32 | Compose, 33 | ) 34 | from pathlib import Path 35 | 36 | from hparam import hparams as hp 37 | 38 | 39 | class MedData_train(torch.utils.data.Dataset): 40 | def __init__(self, images_dir_0, images_dir_1): 41 | 42 | self.subjects = [] 43 | 44 | 45 | images_dir_0 = Path(images_dir_0) 46 | self.image_paths_0 = sorted(images_dir_0.glob(hp.fold_arch)) 47 | 48 | images_dir_1 = Path(images_dir_1) 49 | self.image_paths_1 = sorted(images_dir_1.glob(hp.fold_arch)) 50 | 51 | for (image_path) in zip(self.image_paths_0): 52 | subject = tio.Subject( 53 | source=tio.ScalarImage(image_path), 54 | label= 0, 55 | ) 56 | self.subjects.append(subject) 57 | 58 | for (image_path) in zip(self.image_paths_1): 59 | subject = tio.Subject( 60 | source=tio.ScalarImage(image_path), 61 | label= 1, 62 | ) 63 | self.subjects.append(subject) 64 | 65 | self.transforms = self.transform() 66 | 67 | self.training_set = tio.SubjectsDataset(self.subjects, transform=self.transforms) 68 | 69 | 70 | # one_subject = self.training_set[0] 71 | # one_subject.plot() 72 | 73 | def transform(self): 74 | 75 | 76 | if hp.aug: 77 | training_transform = Compose([ 78 | CropOrPad((hp.crop_or_pad_size), padding_mode='reflect'), 79 | # ToCanonical(), 80 | RandomBiasField(), 81 | ZNormalization(), 82 | RandomNoise(), 83 | RandomFlip(axes=(0,)), 84 | OneOf({ 85 | RandomAffine(): 0.8, 86 | RandomElasticDeformation(): 0.2, 87 | }), 88 | ]) 89 | else: 90 | training_transform = Compose([ 91 | CropOrPad((hp.crop_or_pad_size), padding_mode='reflect'), 92 | ZNormalization(), 93 | ]) 94 | 95 | 96 | return training_transform 97 | 98 | 99 | 100 | 101 | class MedData_test(torch.utils.data.Dataset): 102 | def __init__(self, images_dir_0, images_dir_1): 103 | 104 | self.subjects = [] 105 | 106 | 107 | images_dir_0 = Path(images_dir_0) 108 | self.image_paths_0 = sorted(images_dir_0.glob(hp.fold_arch)) 109 | 110 | images_dir_1 = Path(images_dir_1) 111 | self.image_paths_1 = sorted(images_dir_1.glob(hp.fold_arch)) 112 | 113 | for (image_path) in zip(self.image_paths_0): 114 | subject = tio.Subject( 115 | source=tio.ScalarImage(image_path), 116 | label= 0, 117 | ) 118 | self.subjects.append(subject) 119 | 120 | for (image_path) in zip(self.image_paths_1): 121 | subject = tio.Subject( 122 | source=tio.ScalarImage(image_path), 123 | label= 1, 124 | ) 125 | self.subjects.append(subject) 126 | 127 | 128 | 129 | self.transforms = self.transform() 130 | 131 | self.testing_set = tio.SubjectsDataset(self.subjects, transform=self.transforms) 132 | 133 | 134 | # one_subject = self.training_set[0] 135 | # one_subject.plot() 136 | 137 | def transform(self): 138 | 139 | testing_transform = Compose([ 140 | CropOrPad((hp.crop_or_pad_size), padding_mode='reflect'), 141 | ZNormalization(), 142 | ]) 143 | 144 | 145 | return testing_transform -------------------------------------------------------------------------------- /hparam.py: -------------------------------------------------------------------------------- 1 | class hparams: 2 | 3 | train_or_test = 'train' 4 | output_dir = 'logs/your_program_name' 5 | aug = False 6 | latest_checkpoint_file = 'checkpoint_latest.pt' 7 | total_epochs = 100 8 | epochs_per_checkpoint = 10 9 | batch_size = 2 10 | ckpt = None 11 | init_lr = 0.0002 12 | scheduer_step_size = 20 13 | scheduer_gamma = 0.8 14 | debug = False 15 | mode = '2d' # '2d or '3d' 16 | in_class = 1 17 | out_class = 2 18 | 19 | crop_or_pad_size = 224,224,1 # if 3D: 256,256,256 20 | 21 | fold_arch = '*.png' 22 | 23 | source_train_0_dir = 'train/0' 24 | source_train_1_dir = 'train/1' 25 | source_test_0_dir = 'test/0' 26 | source_test_1_dir = 'test/1' 27 | -------------------------------------------------------------------------------- /loss_function.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | import torch 3 | import torch.nn.functional as F 4 | class FocalLoss(nn.Module): 5 | def __init__(self, alpha=1, gamma=2, logits=True, reduce=True): 6 | super(FocalLoss, self).__init__() 7 | self.alpha = alpha 8 | self.gamma = gamma 9 | self.logits = logits 10 | self.reduce = reduce 11 | 12 | def forward(self, inputs, targets): 13 | if self.logits: 14 | BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduce=False) 15 | else: 16 | BCE_loss = F.binary_cross_entropy(inputs, targets, reduce=False) 17 | pt = torch.exp(-BCE_loss) 18 | F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss 19 | 20 | if self.reduce: 21 | return torch.mean(F_loss) 22 | else: 23 | return F_loss 24 | 25 | class Classification_Loss(nn.Module): 26 | def __init__(self): 27 | super(Classification_Loss, self).__init__() 28 | self.criterionCE = nn.CrossEntropyLoss() 29 | 30 | 31 | def forward(self, model_output, targets, model): 32 | 33 | # torch.empty(3, dtype=torch.long) 34 | # model_output = model_output.long() 35 | # targets = targets.long() 36 | # print(model_output) 37 | # print(F.sigmoid(model_output)) 38 | # print(targets) 39 | # print('kkk') 40 | regularization_loss = 0 41 | for param in model.module.parameters(): 42 | regularization_loss += torch.sum(torch.abs(param)) #+torch.sum(torch.abs(param))**2 43 | # loss = 0.00001 * regularization_loss 44 | loss = 0 45 | 46 | # model_output = F.sigmoid(model_output) 47 | # loss = self.mse_criterion(model_output,targets) 48 | loss += self.criterionCE(model_output,targets) 49 | return loss 50 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['CUDA_VISIBLE_DEVICES'] = '0' 3 | devicess = [0] 4 | from collections import defaultdict 5 | import time 6 | import argparse 7 | import numpy as np 8 | from PIL import Image 9 | import torch 10 | from sklearn import metrics 11 | from torch.utils.data import DataLoader 12 | from torch.utils.tensorboard import SummaryWriter 13 | from torch import nn 14 | from torchvision import transforms 15 | import torch.distributed as dist 16 | import math 17 | import torchio 18 | from torchio.transforms import ( 19 | ZNormalization, 20 | ) 21 | from tqdm import tqdm 22 | from torchvision import utils 23 | from hparam import hparams as hp 24 | from utils.metric import metric 25 | from torch.optim.lr_scheduler import ReduceLROnPlateau,StepLR,CosineAnnealingLR 26 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 27 | 28 | 29 | 30 | source_train_0_dir = hp.source_train_0_dir 31 | source_train_1_dir = hp.source_train_1_dir 32 | 33 | 34 | source_test_0_dir = hp.source_test_0_dir 35 | source_test_1_dir = hp.source_test_1_dir 36 | 37 | 38 | 39 | 40 | def parse_training_args(parser): 41 | """ 42 | Parse commandline arguments. 43 | """ 44 | 45 | parser.add_argument('-o', '--output_dir', type=str, default=hp.output_dir, required=False, help='Directory to save checkpoints') 46 | parser.add_argument('--latest-checkpoint-file', type=str, default=hp.latest_checkpoint_file, help='Store the latest checkpoint in each epoch') 47 | 48 | # training 49 | training = parser.add_argument_group('training setup') 50 | training.add_argument('--epochs', type=int, default=hp.total_epochs, help='Number of total epochs to run') 51 | training.add_argument('--epochs-per-checkpoint', type=int, default=hp.epochs_per_checkpoint, help='Number of epochs per checkpoint') 52 | training.add_argument('--batch', type=int, default=hp.batch_size, help='batch-size') 53 | parser.add_argument( 54 | '-k', 55 | "--ckpt", 56 | type=str, 57 | default=hp.ckpt, 58 | help="path to the checkpoints to resume training", 59 | ) 60 | parser.add_argument("--init-lr", type=float, default=hp.init_lr, help="learning rate") 61 | # TODO 62 | parser.add_argument( 63 | "--local_rank", type=int, default=0, help="local rank for distributed training" 64 | ) 65 | 66 | training.add_argument('--amp-run', action='store_true', help='Enable AMP') 67 | training.add_argument('--cudnn-enabled', default=True, help='Enable cudnn') 68 | training.add_argument('--cudnn-benchmark', default=True, help='Run cudnn benchmark') 69 | training.add_argument('--disable-uniform-initialize-bn-weight', action='store_true', help='disable uniform initialization of batchnorm layer weight') 70 | 71 | return parser 72 | 73 | 74 | 75 | def train(): 76 | 77 | parser = argparse.ArgumentParser(description='PyTorch Medical Segmentation Training') 78 | parser = parse_training_args(parser) 79 | args, _ = parser.parse_known_args() 80 | 81 | args = parser.parse_args() 82 | 83 | 84 | torch.backends.cudnn.deterministic = True 85 | torch.backends.cudnn.enabled = args.cudnn_enabled 86 | torch.backends.cudnn.benchmark = args.cudnn_benchmark 87 | 88 | 89 | from data_function import MedData_train 90 | os.makedirs(args.output_dir, exist_ok=True) 91 | 92 | if hp.mode == '2d': 93 | # from models.two_d.alexnet import alexnet 94 | # model = alexnet(num_classes=2) 95 | 96 | # from models.two_d.densenet import densenet121 97 | # model = densenet121(num_class=2) 98 | 99 | from models.two_d.googlenet import googlenet 100 | model = googlenet(num_class=2) 101 | 102 | # from models.two_d.mobilenet import mobilenet 103 | # model = mobilenet(class_num=2) 104 | 105 | # from models.two_d.nasnet import nasnet 106 | # model = nasnet(class_num=2) 107 | 108 | # from models.two_d.resnet import resnet101 109 | # model = resnet101(num_classes=2) 110 | 111 | # from models.two_d.resnext import resnext101 112 | # model = resnext101(class_names=2) 113 | 114 | # from models.two_d.vggnet import vgg16_bn 115 | # model = vgg16_bn(num_class=2) 116 | 117 | elif hp.mode == '3d': 118 | 119 | from models.three_d.resnet3d import generate_model 120 | model = generate_model(18,n_input_channels=1,n_classes=2) 121 | 122 | from models.three_d.resnext3d import generate_model 123 | model = generate_model(50,n_input_channels=1,n_classes=2) 124 | 125 | from models.three_d.densenet3d import generate_model 126 | model = generate_model(121,n_input_channels=1,num_classes=2) 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | model = torch.nn.DataParallel(model, device_ids=devicess) 136 | optimizer = torch.optim.Adam(model.parameters(), lr=args.init_lr) 137 | 138 | 139 | # scheduler = ReduceLROnPlateau(optimizer, 'min',factor=0.5, patience=20, verbose=True) 140 | scheduler = StepLR(optimizer, step_size=hp.scheduer_step_size, gamma=hp.scheduer_gamma) 141 | # scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=5e-6) 142 | 143 | if args.ckpt is not None: 144 | print("load model:", args.ckpt) 145 | print(os.path.join(args.output_dir, args.latest_checkpoint_file)) 146 | ckpt = torch.load(os.path.join(args.output_dir, args.latest_checkpoint_file), map_location=lambda storage, loc: storage) 147 | 148 | model.load_state_dict(ckpt["model"]) 149 | optimizer.load_state_dict(ckpt["optim"]) 150 | 151 | for state in optimizer.state.values(): 152 | for k, v in state.items(): 153 | if torch.is_tensor(v): 154 | state[k] = v.cuda() 155 | 156 | # scheduler.load_state_dict(ckpt["scheduler"]) 157 | elapsed_epochs = ckpt["epoch"] 158 | else: 159 | elapsed_epochs = 0 160 | 161 | model.cuda() 162 | 163 | from loss_function import Classification_Loss 164 | criterion = Classification_Loss().cuda() 165 | 166 | 167 | writer = SummaryWriter(args.output_dir) 168 | 169 | 170 | 171 | train_dataset = MedData_train(source_train_0_dir,source_train_1_dir) 172 | train_loader = DataLoader(train_dataset.training_set, 173 | batch_size=args.batch, 174 | shuffle=True, 175 | pin_memory=True, 176 | drop_last=True) 177 | 178 | model.train() 179 | 180 | epochs = args.epochs - elapsed_epochs 181 | iteration = elapsed_epochs * len(train_loader) 182 | 183 | 184 | 185 | for epoch in range(1, epochs + 1): 186 | print("epoch:"+str(epoch)) 187 | epoch += elapsed_epochs 188 | 189 | num_iters = 0 190 | 191 | gts = [] 192 | predicts = [] 193 | for i, batch in enumerate(train_loader): 194 | 195 | 196 | if hp.debug: 197 | if i >=1: 198 | break 199 | 200 | print(f"Batch: {i}/{len(train_loader)} epoch {epoch}") 201 | 202 | optimizer.zero_grad() 203 | 204 | 205 | x = batch['source']['data'] 206 | y = batch['label'] 207 | 208 | x = x.type(torch.FloatTensor).cuda() 209 | y = y.type(torch.LongTensor).cuda() 210 | 211 | if hp.mode == '2d': 212 | x = x.squeeze(-1) 213 | x = x[:,:1,:,:] 214 | 215 | outputs = model(x) 216 | 217 | outputs_logit = outputs.argmax(dim=1) 218 | 219 | loss = criterion(outputs, y, model) 220 | 221 | num_iters += 1 222 | loss.backward() 223 | 224 | optimizer.step() 225 | iteration += 1 226 | 227 | 228 | print("loss:"+str(loss.item())) 229 | writer.add_scalar('Training/Loss', loss.item(),iteration) 230 | 231 | predicts.append(outputs_logit.cpu().detach().numpy()) 232 | gts.append(y.cpu().detach().numpy()) 233 | 234 | 235 | predicts = np.concatenate(predicts).flatten().astype(np.int16) 236 | gts = np.concatenate(gts).flatten().astype(np.int16) 237 | 238 | print(metrics.confusion_matrix(predicts, gts)) 239 | acc = metrics.accuracy_score(predicts, gts) 240 | recall = metrics.recall_score(predicts, gts) 241 | f1 = metrics.f1_score(predicts, gts) 242 | writer.add_scalar('Training/acc', acc,epoch) 243 | writer.add_scalar('Training/recall', recall,epoch) 244 | writer.add_scalar('Training/f1', f1,epoch) 245 | 246 | 247 | scheduler.step() 248 | 249 | 250 | # Store latest checkpoint in each epoch 251 | torch.save( 252 | { 253 | "model": model.state_dict(), 254 | "optim": optimizer.state_dict(), 255 | "scheduler":scheduler.state_dict(), 256 | "epoch": epoch, 257 | 258 | }, 259 | os.path.join(args.output_dir, args.latest_checkpoint_file), 260 | ) 261 | 262 | 263 | # Save checkpoint 264 | if epoch % args.epochs_per_checkpoint == 0: 265 | 266 | torch.save( 267 | { 268 | "model": model.state_dict(), 269 | "optim": optimizer.state_dict(), 270 | "epoch": epoch, 271 | }, 272 | os.path.join(args.output_dir, f"checkpoint_{epoch:04d}.pt"), 273 | ) 274 | 275 | writer.close() 276 | 277 | 278 | def test(): 279 | 280 | parser = argparse.ArgumentParser(description='PyTorch Medical Segmentation Testing') 281 | parser = parse_training_args(parser) 282 | args, _ = parser.parse_known_args() 283 | 284 | args = parser.parse_args() 285 | 286 | 287 | torch.backends.cudnn.deterministic = True 288 | torch.backends.cudnn.enabled = args.cudnn_enabled 289 | torch.backends.cudnn.benchmark = args.cudnn_benchmark 290 | 291 | from data_function import MedData_test 292 | 293 | 294 | if hp.mode == '2d': 295 | # from models.two_d.alexnet import alexnet 296 | # model = alexnet(num_classes=2) 297 | 298 | # from models.two_d.densenet import densenet121 299 | # model = densenet121(num_class=2) 300 | 301 | from models.two_d.googlenet import googlenet 302 | model = googlenet(num_class=2) 303 | 304 | # from models.two_d.mobilenet import mobilenet 305 | # model = mobilenet(class_num=2) 306 | 307 | # from models.two_d.nasnet import nasnet 308 | # model = nasnet(class_num=2) 309 | 310 | # from models.two_d.resnet import resnet101 311 | # model = resnet101(num_classes=2) 312 | 313 | # from models.two_d.resnext import resnext101 314 | # model = resnext101(class_names=2) 315 | 316 | # from models.two_d.vggnet import vgg16_bn 317 | # model = vgg16_bn(num_class=2) 318 | 319 | elif hp.mode == '3d': 320 | 321 | from models.three_d.resnet3d import generate_model 322 | model = generate_model(18,n_input_channels=1,n_classes=2) 323 | 324 | from models.three_d.resnext3d import generate_model 325 | model = generate_model(50,n_input_channels=1,n_classes=2) 326 | 327 | from models.three_d.densenet3d import generate_model 328 | model = generate_model(121,n_input_channels=1,num_classes=2) 329 | 330 | 331 | 332 | 333 | model = torch.nn.DataParallel(model, device_ids=devicess,output_device=[1]) 334 | 335 | 336 | print("load model:", args.ckpt) 337 | print(os.path.join(args.output_dir, args.latest_checkpoint_file)) 338 | ckpt = torch.load(os.path.join(args.output_dir, args.latest_checkpoint_file), map_location=lambda storage, loc: storage) 339 | 340 | model.load_state_dict(ckpt["model"]) 341 | 342 | 343 | model.cuda() 344 | 345 | test_dataset = MedData_test(source_test_0_dir,source_test_1_dir) 346 | test_loader = DataLoader(test_dataset.testing_set, 347 | batch_size=args.batch, 348 | shuffle=True, 349 | pin_memory=True, 350 | drop_last=True) 351 | 352 | model.eval() 353 | 354 | predicts = [] 355 | gts = [] 356 | for i, batch in enumerate(test_loader): 357 | 358 | x = batch['source']['data'] 359 | y = batch['label'] 360 | 361 | x = x.type(torch.FloatTensor).cuda() 362 | y = y.type(torch.LongTensor).cuda() 363 | 364 | if hp.mode == '2d': 365 | x = x.squeeze(-1) 366 | x = x[:,:1,:,:] 367 | 368 | outputs = model(x) 369 | 370 | outputs_logit = outputs.argmax(dim=1) 371 | 372 | predicts.append(outputs_logit.cpu().detach().numpy()) 373 | gts.append(y.cpu().detach().numpy()) 374 | 375 | predicts = np.concatenate(predicts).flatten().astype(np.int16) 376 | gts = np.concatenate(gts).flatten().astype(np.int16) 377 | acc = metrics.accuracy_score(predicts, gts) 378 | recall = metrics.recall_score(predicts, gts) 379 | f1 = metrics.f1_score(predicts, gts) 380 | ## log 381 | print("acc:"+str(acc)) 382 | print("recall:"+str(recall)) 383 | print("f1:"+str(f1)) 384 | print(metrics.confusion_matrix(predicts, gts)) 385 | 386 | 387 | 388 | 389 | 390 | 391 | 392 | 393 | 394 | if __name__ == '__main__': 395 | if hp.train_or_test == 'train': 396 | train() 397 | elif hp.train_or_test == 'test': 398 | test() 399 | -------------------------------------------------------------------------------- /models/three_d/densenet3d.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from collections import OrderedDict 7 | 8 | 9 | class _DenseLayer(nn.Sequential): 10 | 11 | def __init__(self, num_input_features, growth_rate, bn_size, drop_rate): 12 | super().__init__() 13 | self.add_module('norm1', nn.BatchNorm3d(num_input_features)) 14 | self.add_module('relu1', nn.ReLU(inplace=True)) 15 | self.add_module( 16 | 'conv1', 17 | nn.Conv3d(num_input_features, 18 | bn_size * growth_rate, 19 | kernel_size=1, 20 | stride=1, 21 | bias=False)) 22 | self.add_module('norm2', nn.BatchNorm3d(bn_size * growth_rate)) 23 | self.add_module('relu2', nn.ReLU(inplace=True)) 24 | self.add_module( 25 | 'conv2', 26 | nn.Conv3d(bn_size * growth_rate, 27 | growth_rate, 28 | kernel_size=3, 29 | stride=1, 30 | padding=1, 31 | bias=False)) 32 | self.drop_rate = drop_rate 33 | 34 | def forward(self, x): 35 | new_features = super().forward(x) 36 | if self.drop_rate > 0: 37 | new_features = F.dropout(new_features, 38 | p=self.drop_rate, 39 | training=self.training) 40 | return torch.cat([x, new_features], 1) 41 | 42 | 43 | class _DenseBlock(nn.Sequential): 44 | 45 | def __init__(self, num_layers, num_input_features, bn_size, growth_rate, 46 | drop_rate): 47 | super().__init__() 48 | for i in range(num_layers): 49 | layer = _DenseLayer(num_input_features + i * growth_rate, 50 | growth_rate, bn_size, drop_rate) 51 | self.add_module('denselayer{}'.format(i + 1), layer) 52 | 53 | 54 | class _Transition(nn.Sequential): 55 | 56 | def __init__(self, num_input_features, num_output_features): 57 | super().__init__() 58 | self.add_module('norm', nn.BatchNorm3d(num_input_features)) 59 | self.add_module('relu', nn.ReLU(inplace=True)) 60 | self.add_module( 61 | 'conv', 62 | nn.Conv3d(num_input_features, 63 | num_output_features, 64 | kernel_size=1, 65 | stride=1, 66 | bias=False)) 67 | self.add_module('pool', nn.AvgPool3d(kernel_size=2, stride=2)) 68 | 69 | 70 | class DenseNet(nn.Module): 71 | """Densenet-BC model class 72 | Args: 73 | growth_rate (int) - how many filters to add each layer (k in paper) 74 | block_config (list of 4 ints) - how many layers in each pooling block 75 | num_init_features (int) - the number of filters to learn in the first convolution layer 76 | bn_size (int) - multiplicative factor for number of bottle neck layers 77 | (i.e. bn_size * k features in the bottleneck layer) 78 | drop_rate (float) - dropout rate after each dense layer 79 | num_classes (int) - number of classification classes 80 | """ 81 | 82 | def __init__(self, 83 | n_input_channels=3, 84 | conv1_t_size=7, 85 | conv1_t_stride=1, 86 | no_max_pool=False, 87 | growth_rate=32, 88 | block_config=(6, 12, 24, 16), 89 | num_init_features=64, 90 | bn_size=4, 91 | drop_rate=0, 92 | num_classes=1000): 93 | 94 | super().__init__() 95 | 96 | # First convolution 97 | self.features = [('conv1', 98 | nn.Conv3d(n_input_channels, 99 | num_init_features, 100 | kernel_size=(conv1_t_size, 7, 7), 101 | stride=(conv1_t_stride, 2, 2), 102 | padding=(conv1_t_size // 2, 3, 3), 103 | bias=False)), 104 | ('norm1', nn.BatchNorm3d(num_init_features)), 105 | ('relu1', nn.ReLU(inplace=True))] 106 | if not no_max_pool: 107 | self.features.append( 108 | ('pool1', nn.MaxPool3d(kernel_size=3, stride=2, padding=1))) 109 | self.features = nn.Sequential(OrderedDict(self.features)) 110 | 111 | # Each denseblock 112 | num_features = num_init_features 113 | for i, num_layers in enumerate(block_config): 114 | block = _DenseBlock(num_layers=num_layers, 115 | num_input_features=num_features, 116 | bn_size=bn_size, 117 | growth_rate=growth_rate, 118 | drop_rate=drop_rate) 119 | self.features.add_module('denseblock{}'.format(i + 1), block) 120 | num_features = num_features + num_layers * growth_rate 121 | if i != len(block_config) - 1: 122 | trans = _Transition(num_input_features=num_features, 123 | num_output_features=num_features // 2) 124 | self.features.add_module('transition{}'.format(i + 1), trans) 125 | num_features = num_features // 2 126 | 127 | # Final batch norm 128 | self.features.add_module('norm5', nn.BatchNorm3d(num_features)) 129 | 130 | for m in self.modules(): 131 | if isinstance(m, nn.Conv3d): 132 | m.weight = nn.init.kaiming_normal(m.weight, mode='fan_out') 133 | elif isinstance(m, nn.BatchNorm3d) or isinstance(m, nn.BatchNorm2d): 134 | m.weight.data.fill_(1) 135 | m.bias.data.zero_() 136 | 137 | # Linear layer 138 | self.classifier = nn.Linear(num_features, num_classes) 139 | 140 | for m in self.modules(): 141 | if isinstance(m, nn.Conv3d): 142 | nn.init.kaiming_normal_(m.weight, 143 | mode='fan_out', 144 | nonlinearity='relu') 145 | elif isinstance(m, nn.BatchNorm3d): 146 | nn.init.constant_(m.weight, 1) 147 | nn.init.constant_(m.bias, 0) 148 | elif isinstance(m, nn.Linear): 149 | nn.init.constant_(m.bias, 0) 150 | 151 | def forward(self, x): 152 | features = self.features(x) 153 | out = F.relu(features, inplace=True) 154 | out = F.adaptive_avg_pool3d(out, 155 | output_size=(1, 1, 156 | 1)).view(features.size(0), -1) 157 | out = self.classifier(out) 158 | return out 159 | 160 | 161 | def generate_model(model_depth, **kwargs): 162 | assert model_depth in [121, 169, 201, 264] 163 | 164 | if model_depth == 121: 165 | model = DenseNet(num_init_features=64, 166 | growth_rate=32, 167 | block_config=(6, 12, 24, 16), 168 | **kwargs) 169 | elif model_depth == 169: 170 | model = DenseNet(num_init_features=64, 171 | growth_rate=32, 172 | block_config=(6, 12, 32, 32), 173 | **kwargs) 174 | elif model_depth == 201: 175 | model = DenseNet(num_init_features=64, 176 | growth_rate=32, 177 | block_config=(6, 12, 48, 32), 178 | **kwargs) 179 | elif model_depth == 264: 180 | model = DenseNet(num_init_features=64, 181 | growth_rate=32, 182 | block_config=(6, 12, 64, 48), 183 | **kwargs) 184 | 185 | return model 186 | 187 | 188 | if __name__ == "__main__": 189 | 190 | 191 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 192 | image_size = 64 193 | x = torch.Tensor(1, 1, image_size, image_size, image_size) 194 | x = x.to(device) 195 | print("x size: {}".format(x.size())) 196 | 197 | model = generate_model(121,n_input_channels=1,num_classes=2).to(device) 198 | 199 | 200 | out1 = model(x) 201 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/three_d/resnet3d.py: -------------------------------------------------------------------------------- 1 | import math 2 | from functools import partial 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | 9 | def get_inplanes(): 10 | return [64, 128, 256, 512] 11 | 12 | 13 | def conv3x3x3(in_planes, out_planes, stride=1): 14 | return nn.Conv3d(in_planes, 15 | out_planes, 16 | kernel_size=3, 17 | stride=stride, 18 | padding=1, 19 | bias=False) 20 | 21 | 22 | def conv1x1x1(in_planes, out_planes, stride=1): 23 | return nn.Conv3d(in_planes, 24 | out_planes, 25 | kernel_size=1, 26 | stride=stride, 27 | bias=False) 28 | 29 | 30 | class BasicBlock(nn.Module): 31 | expansion = 1 32 | 33 | def __init__(self, in_planes, planes, stride=1, downsample=None): 34 | super().__init__() 35 | 36 | self.conv1 = conv3x3x3(in_planes, planes, stride) 37 | self.bn1 = nn.BatchNorm3d(planes) 38 | self.relu = nn.ReLU(inplace=True) 39 | self.conv2 = conv3x3x3(planes, planes) 40 | self.bn2 = nn.BatchNorm3d(planes) 41 | self.downsample = downsample 42 | self.stride = stride 43 | 44 | def forward(self, x): 45 | residual = x 46 | 47 | out = self.conv1(x) 48 | out = self.bn1(out) 49 | out = self.relu(out) 50 | 51 | out = self.conv2(out) 52 | out = self.bn2(out) 53 | 54 | if self.downsample is not None: 55 | residual = self.downsample(x) 56 | 57 | out += residual 58 | out = self.relu(out) 59 | 60 | return out 61 | 62 | 63 | class Bottleneck(nn.Module): 64 | expansion = 4 65 | 66 | def __init__(self, in_planes, planes, stride=1, downsample=None): 67 | super().__init__() 68 | 69 | self.conv1 = conv1x1x1(in_planes, planes) 70 | self.bn1 = nn.BatchNorm3d(planes) 71 | self.conv2 = conv3x3x3(planes, planes, stride) 72 | self.bn2 = nn.BatchNorm3d(planes) 73 | self.conv3 = conv1x1x1(planes, planes * self.expansion) 74 | self.bn3 = nn.BatchNorm3d(planes * self.expansion) 75 | self.relu = nn.ReLU(inplace=True) 76 | self.downsample = downsample 77 | self.stride = stride 78 | 79 | def forward(self, x): 80 | residual = x 81 | 82 | out = self.conv1(x) 83 | out = self.bn1(out) 84 | out = self.relu(out) 85 | 86 | out = self.conv2(out) 87 | out = self.bn2(out) 88 | out = self.relu(out) 89 | 90 | out = self.conv3(out) 91 | out = self.bn3(out) 92 | 93 | if self.downsample is not None: 94 | residual = self.downsample(x) 95 | 96 | out += residual 97 | out = self.relu(out) 98 | 99 | return out 100 | 101 | 102 | class ResNet(nn.Module): 103 | 104 | def __init__(self, 105 | block, 106 | layers, 107 | block_inplanes, 108 | n_input_channels=3, 109 | conv1_t_size=7, 110 | conv1_t_stride=1, 111 | no_max_pool=False, 112 | shortcut_type='B', 113 | widen_factor=1.0, 114 | n_classes=400): 115 | super().__init__() 116 | 117 | block_inplanes = [int(x * widen_factor) for x in block_inplanes] 118 | 119 | self.in_planes = block_inplanes[0] 120 | self.no_max_pool = no_max_pool 121 | 122 | self.conv1 = nn.Conv3d(n_input_channels, 123 | self.in_planes, 124 | kernel_size=(conv1_t_size, 7, 7), 125 | stride=(conv1_t_stride, 2, 2), 126 | padding=(conv1_t_size // 2, 3, 3), 127 | bias=False) 128 | self.bn1 = nn.BatchNorm3d(self.in_planes) 129 | self.relu = nn.ReLU(inplace=True) 130 | self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1) 131 | self.layer1 = self._make_layer(block, block_inplanes[0], layers[0], 132 | shortcut_type) 133 | self.layer2 = self._make_layer(block, 134 | block_inplanes[1], 135 | layers[1], 136 | shortcut_type, 137 | stride=2) 138 | self.layer3 = self._make_layer(block, 139 | block_inplanes[2], 140 | layers[2], 141 | shortcut_type, 142 | stride=2) 143 | self.layer4 = self._make_layer(block, 144 | block_inplanes[3], 145 | layers[3], 146 | shortcut_type, 147 | stride=2) 148 | 149 | self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1)) 150 | self.fc = nn.Linear(block_inplanes[3] * block.expansion, n_classes) 151 | 152 | for m in self.modules(): 153 | if isinstance(m, nn.Conv3d): 154 | nn.init.kaiming_normal_(m.weight, 155 | mode='fan_out', 156 | nonlinearity='relu') 157 | elif isinstance(m, nn.BatchNorm3d): 158 | nn.init.constant_(m.weight, 1) 159 | nn.init.constant_(m.bias, 0) 160 | 161 | def _downsample_basic_block(self, x, planes, stride): 162 | out = F.avg_pool3d(x, kernel_size=1, stride=stride) 163 | zero_pads = torch.zeros(out.size(0), planes - out.size(1), out.size(2), 164 | out.size(3), out.size(4)) 165 | if isinstance(out.data, torch.cuda.FloatTensor): 166 | zero_pads = zero_pads.cuda() 167 | 168 | out = torch.cat([out.data, zero_pads], dim=1) 169 | 170 | return out 171 | 172 | def _make_layer(self, block, planes, blocks, shortcut_type, stride=1): 173 | downsample = None 174 | if stride != 1 or self.in_planes != planes * block.expansion: 175 | if shortcut_type == 'A': 176 | downsample = partial(self._downsample_basic_block, 177 | planes=planes * block.expansion, 178 | stride=stride) 179 | else: 180 | downsample = nn.Sequential( 181 | conv1x1x1(self.in_planes, planes * block.expansion, stride), 182 | nn.BatchNorm3d(planes * block.expansion)) 183 | 184 | layers = [] 185 | layers.append( 186 | block(in_planes=self.in_planes, 187 | planes=planes, 188 | stride=stride, 189 | downsample=downsample)) 190 | self.in_planes = planes * block.expansion 191 | for i in range(1, blocks): 192 | layers.append(block(self.in_planes, planes)) 193 | 194 | return nn.Sequential(*layers) 195 | 196 | def forward(self, x): 197 | x = self.conv1(x) 198 | x = self.bn1(x) 199 | x = self.relu(x) 200 | if not self.no_max_pool: 201 | x = self.maxpool(x) 202 | 203 | x = self.layer1(x) 204 | x = self.layer2(x) 205 | x = self.layer3(x) 206 | x = self.layer4(x) 207 | 208 | x = self.avgpool(x) 209 | 210 | x = x.view(x.size(0), -1) 211 | x = self.fc(x) 212 | 213 | return x 214 | 215 | 216 | def generate_model(model_depth, **kwargs): 217 | assert model_depth in [10, 18, 34, 50, 101, 152, 200] 218 | 219 | if model_depth == 10: 220 | model = ResNet(BasicBlock, [1, 1, 1, 1], get_inplanes(), **kwargs) 221 | elif model_depth == 18: 222 | model = ResNet(BasicBlock, [2, 2, 2, 2], get_inplanes(), **kwargs) 223 | elif model_depth == 34: 224 | model = ResNet(BasicBlock, [3, 4, 6, 3], get_inplanes(), **kwargs) 225 | elif model_depth == 50: 226 | model = ResNet(Bottleneck, [3, 4, 6, 3], get_inplanes(), **kwargs) 227 | elif model_depth == 101: 228 | model = ResNet(Bottleneck, [3, 4, 23, 3], get_inplanes(), **kwargs) 229 | elif model_depth == 152: 230 | model = ResNet(Bottleneck, [3, 8, 36, 3], get_inplanes(), **kwargs) 231 | elif model_depth == 200: 232 | model = ResNet(Bottleneck, [3, 24, 36, 3], get_inplanes(), **kwargs) 233 | 234 | return model 235 | 236 | 237 | if __name__ == "__main__": 238 | 239 | 240 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 241 | image_size = 64 242 | x = torch.Tensor(1, 1, image_size, image_size, image_size) 243 | x = x.to(device) 244 | print("x size: {}".format(x.size())) 245 | 246 | model = generate_model(18,n_input_channels=1,n_classes=2).to(device) 247 | 248 | 249 | out1 = model(x) 250 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/three_d/resnext3d.py: -------------------------------------------------------------------------------- 1 | import math 2 | from functools import partial, partialmethod 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | 9 | def conv3x3x3(in_planes, out_planes, stride=1): 10 | return nn.Conv3d(in_planes, 11 | out_planes, 12 | kernel_size=3, 13 | stride=stride, 14 | padding=1, 15 | bias=False) 16 | 17 | 18 | def conv1x1x1(in_planes, out_planes, stride=1): 19 | return nn.Conv3d(in_planes, 20 | out_planes, 21 | kernel_size=1, 22 | stride=stride, 23 | bias=False) 24 | 25 | 26 | class Bottleneck(nn.Module): 27 | expansion = 4 28 | 29 | def __init__(self, in_planes, planes, stride=1, downsample=None): 30 | super().__init__() 31 | 32 | self.conv1 = conv1x1x1(in_planes, planes) 33 | self.bn1 = nn.BatchNorm3d(planes) 34 | self.conv2 = conv3x3x3(planes, planes, stride) 35 | self.bn2 = nn.BatchNorm3d(planes) 36 | self.conv3 = conv1x1x1(planes, planes * self.expansion) 37 | self.bn3 = nn.BatchNorm3d(planes * self.expansion) 38 | self.relu = nn.ReLU(inplace=True) 39 | self.downsample = downsample 40 | self.stride = stride 41 | 42 | def forward(self, x): 43 | residual = x 44 | 45 | out = self.conv1(x) 46 | out = self.bn1(out) 47 | out = self.relu(out) 48 | 49 | out = self.conv2(out) 50 | out = self.bn2(out) 51 | out = self.relu(out) 52 | 53 | out = self.conv3(out) 54 | out = self.bn3(out) 55 | 56 | if self.downsample is not None: 57 | residual = self.downsample(x) 58 | 59 | out += residual 60 | out = self.relu(out) 61 | 62 | return out 63 | 64 | 65 | class ResNet(nn.Module): 66 | 67 | def __init__(self, 68 | block, 69 | layers, 70 | block_inplanes, 71 | n_input_channels=3, 72 | conv1_t_size=7, 73 | conv1_t_stride=1, 74 | no_max_pool=False, 75 | shortcut_type='B', 76 | widen_factor=1.0, 77 | n_classes=400): 78 | super().__init__() 79 | 80 | block_inplanes = [int(x * widen_factor) for x in block_inplanes] 81 | 82 | self.in_planes = block_inplanes[0] 83 | self.no_max_pool = no_max_pool 84 | 85 | self.conv1 = nn.Conv3d(n_input_channels, 86 | self.in_planes, 87 | kernel_size=(conv1_t_size, 7, 7), 88 | stride=(conv1_t_stride, 2, 2), 89 | padding=(conv1_t_size // 2, 3, 3), 90 | bias=False) 91 | self.bn1 = nn.BatchNorm3d(self.in_planes) 92 | self.relu = nn.ReLU(inplace=True) 93 | self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1) 94 | self.layer1 = self._make_layer(block, block_inplanes[0], layers[0], 95 | shortcut_type) 96 | self.layer2 = self._make_layer(block, 97 | block_inplanes[1], 98 | layers[1], 99 | shortcut_type, 100 | stride=2) 101 | self.layer3 = self._make_layer(block, 102 | block_inplanes[2], 103 | layers[2], 104 | shortcut_type, 105 | stride=2) 106 | self.layer4 = self._make_layer(block, 107 | block_inplanes[3], 108 | layers[3], 109 | shortcut_type, 110 | stride=2) 111 | 112 | self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1)) 113 | self.fc = nn.Linear(block_inplanes[3] * block.expansion, n_classes) 114 | 115 | for m in self.modules(): 116 | if isinstance(m, nn.Conv3d): 117 | nn.init.kaiming_normal_(m.weight, 118 | mode='fan_out', 119 | nonlinearity='relu') 120 | elif isinstance(m, nn.BatchNorm3d): 121 | nn.init.constant_(m.weight, 1) 122 | nn.init.constant_(m.bias, 0) 123 | 124 | def _downsample_basic_block(self, x, planes, stride): 125 | out = F.avg_pool3d(x, kernel_size=1, stride=stride) 126 | zero_pads = torch.zeros(out.size(0), planes - out.size(1), out.size(2), 127 | out.size(3), out.size(4)) 128 | if isinstance(out.data, torch.cuda.FloatTensor): 129 | zero_pads = zero_pads.cuda() 130 | 131 | out = torch.cat([out.data, zero_pads], dim=1) 132 | 133 | return out 134 | 135 | def _make_layer(self, block, planes, blocks, shortcut_type, stride=1): 136 | downsample = None 137 | if stride != 1 or self.in_planes != planes * block.expansion: 138 | if shortcut_type == 'A': 139 | downsample = partial(self._downsample_basic_block, 140 | planes=planes * block.expansion, 141 | stride=stride) 142 | else: 143 | downsample = nn.Sequential( 144 | conv1x1x1(self.in_planes, planes * block.expansion, stride), 145 | nn.BatchNorm3d(planes * block.expansion)) 146 | 147 | layers = [] 148 | layers.append( 149 | block(in_planes=self.in_planes, 150 | planes=planes, 151 | stride=stride, 152 | downsample=downsample)) 153 | self.in_planes = planes * block.expansion 154 | for i in range(1, blocks): 155 | layers.append(block(self.in_planes, planes)) 156 | 157 | return nn.Sequential(*layers) 158 | 159 | def forward(self, x): 160 | x = self.conv1(x) 161 | x = self.bn1(x) 162 | x = self.relu(x) 163 | if not self.no_max_pool: 164 | x = self.maxpool(x) 165 | 166 | x = self.layer1(x) 167 | x = self.layer2(x) 168 | x = self.layer3(x) 169 | x = self.layer4(x) 170 | 171 | x = self.avgpool(x) 172 | 173 | x = x.view(x.size(0), -1) 174 | x = self.fc(x) 175 | 176 | return x 177 | 178 | 179 | 180 | def partialclass(cls, *args, **kwargs): 181 | 182 | class PartialClass(cls): 183 | __init__ = partialmethod(cls.__init__, *args, **kwargs) 184 | 185 | return PartialClass 186 | 187 | 188 | def get_inplanes(): 189 | return [128, 256, 512, 1024] 190 | 191 | 192 | class ResNeXtBottleneck(Bottleneck): 193 | expansion = 2 194 | 195 | def __init__(self, in_planes, planes, cardinality, stride=1, 196 | downsample=None): 197 | super().__init__(in_planes, planes, stride, downsample) 198 | 199 | mid_planes = cardinality * planes // 32 200 | self.conv1 = conv1x1x1(in_planes, mid_planes) 201 | self.bn1 = nn.BatchNorm3d(mid_planes) 202 | self.conv2 = nn.Conv3d(mid_planes, 203 | mid_planes, 204 | kernel_size=3, 205 | stride=stride, 206 | padding=1, 207 | groups=cardinality, 208 | bias=False) 209 | self.bn2 = nn.BatchNorm3d(mid_planes) 210 | self.conv3 = conv1x1x1(mid_planes, planes * self.expansion) 211 | 212 | 213 | class ResNeXt(ResNet): 214 | 215 | def __init__(self, 216 | block, 217 | layers, 218 | block_inplanes, 219 | n_input_channels=3, 220 | conv1_t_size=7, 221 | conv1_t_stride=1, 222 | no_max_pool=False, 223 | shortcut_type='B', 224 | cardinality=64, 225 | n_classes=400): 226 | block = partialclass(block, cardinality=cardinality) 227 | super().__init__(block, layers, block_inplanes, n_input_channels, 228 | conv1_t_size, conv1_t_stride, no_max_pool, 229 | shortcut_type, n_classes) 230 | 231 | self.fc = nn.Linear(cardinality * 16 * block.expansion, n_classes) 232 | 233 | 234 | def generate_model(model_depth, **kwargs): 235 | assert model_depth in [50, 101, 152, 200] 236 | 237 | if model_depth == 50: 238 | model = ResNeXt(ResNeXtBottleneck, [3, 4, 6, 3], get_inplanes(), 239 | **kwargs) 240 | elif model_depth == 101: 241 | model = ResNeXt(ResNeXtBottleneck, [3, 4, 23, 3], get_inplanes(), 242 | **kwargs) 243 | elif model_depth == 152: 244 | model = ResNeXt(ResNeXtBottleneck, [3, 8, 36, 3], get_inplanes(), 245 | **kwargs) 246 | elif model_depth == 200: 247 | model = ResNeXt(ResNeXtBottleneck, [3, 24, 36, 3], get_inplanes(), 248 | **kwargs) 249 | 250 | return model 251 | 252 | 253 | if __name__ == "__main__": 254 | 255 | 256 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 257 | image_size = 64 258 | x = torch.Tensor(1, 1, image_size, image_size, image_size) 259 | x = x.to(device) 260 | print("x size: {}".format(x.size())) 261 | 262 | model = generate_model(50,n_input_channels=1,n_classes=1).to(device) 263 | 264 | 265 | out1 = model(x) 266 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/two_d/alexnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | import torch 4 | 5 | class AlexNet(nn.Module): 6 | 7 | def __init__(self, num_classes=1000): 8 | super(AlexNet, self).__init__() 9 | self.features = nn.Sequential( 10 | nn.Conv2d(1, 64, kernel_size=11, stride=4, padding=2), 11 | nn.ReLU(inplace=True), 12 | nn.MaxPool2d(kernel_size=3, stride=2), 13 | nn.Conv2d(64, 192, kernel_size=5, padding=2), 14 | nn.ReLU(inplace=True), 15 | nn.MaxPool2d(kernel_size=3, stride=2), 16 | nn.Conv2d(192, 384, kernel_size=3, padding=1), 17 | nn.ReLU(inplace=True), 18 | nn.Conv2d(384, 256, kernel_size=3, padding=1), 19 | nn.ReLU(inplace=True), 20 | nn.Conv2d(256, 256, kernel_size=3, padding=1), 21 | nn.ReLU(inplace=True), 22 | nn.MaxPool2d(kernel_size=3, stride=2), 23 | ) 24 | self.classifier = nn.Sequential( 25 | nn.Dropout(), 26 | nn.Linear(256 * 6 * 6, 4096), 27 | nn.ReLU(inplace=True), 28 | nn.Dropout(), 29 | nn.Linear(4096, 4096), 30 | nn.ReLU(inplace=True), 31 | nn.Linear(4096, num_classes), 32 | ) 33 | 34 | def forward(self, x): 35 | x = self.features(x) 36 | x = x.view(x.size(0), 256 * 6 * 6) 37 | x = self.classifier(x) 38 | return x 39 | 40 | 41 | def alexnet(num_classes): 42 | return AlexNet(num_classes) 43 | 44 | 45 | if __name__ == "__main__": 46 | 47 | 48 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 49 | image_size = 224 50 | x = torch.Tensor(1, 1, image_size, image_size) 51 | x = x.to(device) 52 | print("x size: {}".format(x.size())) 53 | 54 | model = alexnet(num_classes=2).to(device) 55 | 56 | 57 | out1 = model(x) 58 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/two_d/densenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | 6 | #"""Bottleneck layers. Although each layer only produces k 7 | #output feature-maps, it typically has many more inputs. It 8 | #has been noted in [37, 11] that a 1×1 convolution can be in- 9 | #troduced as bottleneck layer before each 3×3 convolution 10 | #to reduce the number of input feature-maps, and thus to 11 | #improve computational efficiency.""" 12 | class Bottleneck(nn.Module): 13 | def __init__(self, in_channels, growth_rate): 14 | super().__init__() 15 | #"""In our experiments, we let each 1×1 convolution 16 | #produce 4k feature-maps.""" 17 | inner_channel = 4 * growth_rate 18 | 19 | #"""We find this design especially effective for DenseNet and 20 | #we refer to our network with such a bottleneck layer, i.e., 21 | #to the BN-ReLU-Conv(1×1)-BN-ReLU-Conv(3×3) version of H ` , 22 | #as DenseNet-B.""" 23 | self.bottle_neck = nn.Sequential( 24 | nn.BatchNorm2d(in_channels), 25 | nn.ReLU(inplace=True), 26 | nn.Conv2d(in_channels, inner_channel, kernel_size=1, bias=False), 27 | nn.BatchNorm2d(inner_channel), 28 | nn.ReLU(inplace=True), 29 | nn.Conv2d(inner_channel, growth_rate, kernel_size=3, padding=1, bias=False) 30 | ) 31 | 32 | def forward(self, x): 33 | return torch.cat([x, self.bottle_neck(x)], 1) 34 | 35 | #"""We refer to layers between blocks as transition 36 | #layers, which do convolution and pooling.""" 37 | class Transition(nn.Module): 38 | def __init__(self, in_channels, out_channels): 39 | super().__init__() 40 | #"""The transition layers used in our experiments 41 | #consist of a batch normalization layer and an 1×1 42 | #convolutional layer followed by a 2×2 average pooling 43 | #layer""". 44 | self.down_sample = nn.Sequential( 45 | nn.BatchNorm2d(in_channels), 46 | nn.Conv2d(in_channels, out_channels, 1, bias=False), 47 | nn.AvgPool2d(2, stride=2) 48 | ) 49 | 50 | def forward(self, x): 51 | return self.down_sample(x) 52 | 53 | #DesneNet-BC 54 | #B stands for bottleneck layer(BN-RELU-CONV(1x1)-BN-RELU-CONV(3x3)) 55 | #C stands for compression factor(0<=theta<=1) 56 | class DenseNet(nn.Module): 57 | def __init__(self, block, nblocks, growth_rate=12, reduction=0.5, num_class=100): 58 | super().__init__() 59 | self.growth_rate = growth_rate 60 | 61 | #"""Before entering the first dense block, a convolution 62 | #with 16 (or twice the growth rate for DenseNet-BC) 63 | #output channels is performed on the input images.""" 64 | inner_channels = 2 * growth_rate 65 | 66 | #For convolutional layers with kernel size 3×3, each 67 | #side of the inputs is zero-padded by one pixel to keep 68 | #the feature-map size fixed. 69 | self.conv1 = nn.Conv2d(1, inner_channels, kernel_size=3, padding=1, bias=False) 70 | 71 | self.features = nn.Sequential() 72 | 73 | for index in range(len(nblocks) - 1): 74 | self.features.add_module("dense_block_layer_{}".format(index), self._make_dense_layers(block, inner_channels, nblocks[index])) 75 | inner_channels += growth_rate * nblocks[index] 76 | 77 | #"""If a dense block contains m feature-maps, we let the 78 | #following transition layer generate θm output feature- 79 | #maps, where 0 < θ ≤ 1 is referred to as the compression 80 | #fac-tor. 81 | out_channels = int(reduction * inner_channels) # int() will automatic floor the value 82 | self.features.add_module("transition_layer_{}".format(index), Transition(inner_channels, out_channels)) 83 | inner_channels = out_channels 84 | 85 | self.features.add_module("dense_block{}".format(len(nblocks) - 1), self._make_dense_layers(block, inner_channels, nblocks[len(nblocks)-1])) 86 | inner_channels += growth_rate * nblocks[len(nblocks) - 1] 87 | self.features.add_module('bn', nn.BatchNorm2d(inner_channels)) 88 | self.features.add_module('relu', nn.ReLU(inplace=True)) 89 | 90 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 91 | 92 | self.linear = nn.Linear(inner_channels, num_class) 93 | 94 | def forward(self, x): 95 | output = self.conv1(x) 96 | output = self.features(output) 97 | output = self.avgpool(output) 98 | output = output.view(output.size()[0], -1) 99 | output = self.linear(output) 100 | return output 101 | 102 | def _make_dense_layers(self, block, in_channels, nblocks): 103 | dense_block = nn.Sequential() 104 | for index in range(nblocks): 105 | dense_block.add_module('bottle_neck_layer_{}'.format(index), block(in_channels, self.growth_rate)) 106 | in_channels += self.growth_rate 107 | return dense_block 108 | 109 | def densenet121(num_class): 110 | return DenseNet(Bottleneck, [6,12,24,16], growth_rate=32, num_class=num_class) 111 | 112 | def densenet169(num_class): 113 | return DenseNet(Bottleneck, [6,12,32,32], growth_rate=32, num_class=num_class) 114 | 115 | def densenet201(num_class): 116 | return DenseNet(Bottleneck, [6,12,48,32], growth_rate=32, num_class=num_class) 117 | 118 | def densenet161(num_class): 119 | return DenseNet(Bottleneck, [6,12,36,24], growth_rate=48, num_class=num_class) 120 | 121 | 122 | if __name__ == "__main__": 123 | 124 | 125 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 126 | image_size = 256 127 | x = torch.Tensor(1, 1, image_size, image_size) 128 | x = x.to(device) 129 | print("x size: {}".format(x.size())) 130 | 131 | model = densenet161(num_class=2).to(device) 132 | 133 | 134 | out1 = model(x) 135 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/two_d/googlenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class Inception(nn.Module): 5 | def __init__(self, input_channels, n1x1, n3x3_reduce, n3x3, n5x5_reduce, n5x5, pool_proj): 6 | super().__init__() 7 | 8 | #1x1conv branch 9 | self.b1 = nn.Sequential( 10 | nn.Conv2d(input_channels, n1x1, kernel_size=1), 11 | nn.BatchNorm2d(n1x1), 12 | nn.ReLU(inplace=True) 13 | ) 14 | 15 | #1x1conv -> 3x3conv branch 16 | self.b2 = nn.Sequential( 17 | nn.Conv2d(input_channels, n3x3_reduce, kernel_size=1), 18 | nn.BatchNorm2d(n3x3_reduce), 19 | nn.ReLU(inplace=True), 20 | nn.Conv2d(n3x3_reduce, n3x3, kernel_size=3, padding=1), 21 | nn.BatchNorm2d(n3x3), 22 | nn.ReLU(inplace=True) 23 | ) 24 | 25 | #1x1conv -> 5x5conv branch 26 | #we use 2 3x3 conv filters stacked instead 27 | #of 1 5x5 filters to obtain the same receptive 28 | #field with fewer parameters 29 | self.b3 = nn.Sequential( 30 | nn.Conv2d(input_channels, n5x5_reduce, kernel_size=1), 31 | nn.BatchNorm2d(n5x5_reduce), 32 | nn.ReLU(inplace=True), 33 | nn.Conv2d(n5x5_reduce, n5x5, kernel_size=3, padding=1), 34 | nn.BatchNorm2d(n5x5, n5x5), 35 | nn.ReLU(inplace=True), 36 | nn.Conv2d(n5x5, n5x5, kernel_size=3, padding=1), 37 | nn.BatchNorm2d(n5x5), 38 | nn.ReLU(inplace=True) 39 | ) 40 | 41 | #3x3pooling -> 1x1conv 42 | #same conv 43 | self.b4 = nn.Sequential( 44 | nn.MaxPool2d(3, stride=1, padding=1), 45 | nn.Conv2d(input_channels, pool_proj, kernel_size=1), 46 | nn.BatchNorm2d(pool_proj), 47 | nn.ReLU(inplace=True) 48 | ) 49 | 50 | def forward(self, x): 51 | return torch.cat([self.b1(x), self.b2(x), self.b3(x), self.b4(x)], dim=1) 52 | 53 | 54 | class GoogleNet(nn.Module): 55 | 56 | def __init__(self, num_class=100): 57 | super().__init__() 58 | self.prelayer = nn.Sequential( 59 | nn.Conv2d(1, 64, kernel_size=3, padding=1, bias=False), 60 | nn.BatchNorm2d(64), 61 | nn.ReLU(inplace=True), 62 | nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False), 63 | nn.BatchNorm2d(64), 64 | nn.ReLU(inplace=True), 65 | nn.Conv2d(64, 192, kernel_size=3, padding=1, bias=False), 66 | nn.BatchNorm2d(192), 67 | nn.ReLU(inplace=True), 68 | ) 69 | 70 | #although we only use 1 conv layer as prelayer, 71 | #we still use name a3, b3....... 72 | self.a3 = Inception(192, 64, 96, 128, 16, 32, 32) 73 | self.b3 = Inception(256, 128, 128, 192, 32, 96, 64) 74 | 75 | ##"""In general, an Inception network is a network consisting of 76 | ##modules of the above type stacked upon each other, with occasional 77 | ##max-pooling layers with stride 2 to halve the resolution of the 78 | ##grid""" 79 | self.maxpool = nn.MaxPool2d(3, stride=2, padding=1) 80 | 81 | self.a4 = Inception(480, 192, 96, 208, 16, 48, 64) 82 | self.b4 = Inception(512, 160, 112, 224, 24, 64, 64) 83 | self.c4 = Inception(512, 128, 128, 256, 24, 64, 64) 84 | self.d4 = Inception(512, 112, 144, 288, 32, 64, 64) 85 | self.e4 = Inception(528, 256, 160, 320, 32, 128, 128) 86 | 87 | self.a5 = Inception(832, 256, 160, 320, 32, 128, 128) 88 | self.b5 = Inception(832, 384, 192, 384, 48, 128, 128) 89 | 90 | #input feature size: 8*8*1024 91 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 92 | self.dropout = nn.Dropout2d(p=0.4) 93 | self.linear = nn.Linear(1024, num_class) 94 | 95 | def forward(self, x): 96 | x = self.prelayer(x) 97 | x = self.maxpool(x) 98 | x = self.a3(x) 99 | x = self.b3(x) 100 | 101 | x = self.maxpool(x) 102 | 103 | x = self.a4(x) 104 | x = self.b4(x) 105 | x = self.c4(x) 106 | x = self.d4(x) 107 | x = self.e4(x) 108 | 109 | x = self.maxpool(x) 110 | 111 | x = self.a5(x) 112 | x = self.b5(x) 113 | 114 | #"""It was found that a move from fully connected layers to 115 | #average pooling improved the top-1 accuracy by about 0.6%, 116 | #however the use of dropout remained essential even after 117 | #removing the fully connected layers.""" 118 | x = self.avgpool(x) 119 | x = self.dropout(x) 120 | x = x.view(x.size()[0], -1) 121 | x = self.linear(x) 122 | 123 | return x 124 | 125 | def googlenet(num_class): 126 | return GoogleNet(num_class=num_class) 127 | 128 | 129 | 130 | if __name__ == "__main__": 131 | 132 | 133 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 134 | image_size = 256 135 | x = torch.Tensor(1, 1, image_size, image_size) 136 | x = x.to(device) 137 | print("x size: {}".format(x.size())) 138 | 139 | model = googlenet(num_class=2).to(device) 140 | 141 | 142 | out1 = model(x) 143 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/two_d/mobilenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class DepthSeperabelConv2d(nn.Module): 6 | 7 | def __init__(self, input_channels, output_channels, kernel_size, **kwargs): 8 | super().__init__() 9 | self.depthwise = nn.Sequential( 10 | nn.Conv2d( 11 | input_channels, 12 | input_channels, 13 | kernel_size, 14 | groups=input_channels, 15 | **kwargs), 16 | nn.BatchNorm2d(input_channels), 17 | nn.ReLU(inplace=True) 18 | ) 19 | 20 | self.pointwise = nn.Sequential( 21 | nn.Conv2d(input_channels, output_channels, 1), 22 | nn.BatchNorm2d(output_channels), 23 | nn.ReLU(inplace=True) 24 | ) 25 | 26 | def forward(self, x): 27 | x = self.depthwise(x) 28 | x = self.pointwise(x) 29 | 30 | return x 31 | 32 | 33 | class BasicConv2d(nn.Module): 34 | 35 | def __init__(self, input_channels, output_channels, kernel_size, **kwargs): 36 | 37 | super().__init__() 38 | self.conv = nn.Conv2d( 39 | input_channels, output_channels, kernel_size, **kwargs) 40 | self.bn = nn.BatchNorm2d(output_channels) 41 | self.relu = nn.ReLU(inplace=True) 42 | 43 | def forward(self, x): 44 | x = self.conv(x) 45 | x = self.bn(x) 46 | x = self.relu(x) 47 | 48 | return x 49 | 50 | 51 | class MobileNet(nn.Module): 52 | 53 | """ 54 | Args: 55 | width multipler: The role of the width multiplier α is to thin 56 | a network uniformly at each layer. For a given 57 | layer and width multiplier α, the number of 58 | input channels M becomes αM and the number of 59 | output channels N becomes αN. 60 | """ 61 | 62 | def __init__(self, width_multiplier=1, class_num=100): 63 | super().__init__() 64 | 65 | alpha = width_multiplier 66 | self.stem = nn.Sequential( 67 | BasicConv2d(1, int(32 * alpha), 3, padding=1, bias=False), 68 | DepthSeperabelConv2d( 69 | int(32 * alpha), 70 | int(64 * alpha), 71 | 3, 72 | padding=1, 73 | bias=False 74 | ) 75 | ) 76 | 77 | #downsample 78 | self.conv1 = nn.Sequential( 79 | DepthSeperabelConv2d( 80 | int(64 * alpha), 81 | int(128 * alpha), 82 | 3, 83 | stride=2, 84 | padding=1, 85 | bias=False 86 | ), 87 | DepthSeperabelConv2d( 88 | int(128 * alpha), 89 | int(128 * alpha), 90 | 3, 91 | padding=1, 92 | bias=False 93 | ) 94 | ) 95 | 96 | #downsample 97 | self.conv2 = nn.Sequential( 98 | DepthSeperabelConv2d( 99 | int(128 * alpha), 100 | int(256 * alpha), 101 | 3, 102 | stride=2, 103 | padding=1, 104 | bias=False 105 | ), 106 | DepthSeperabelConv2d( 107 | int(256 * alpha), 108 | int(256 * alpha), 109 | 3, 110 | padding=1, 111 | bias=False 112 | ) 113 | ) 114 | 115 | #downsample 116 | self.conv3 = nn.Sequential( 117 | DepthSeperabelConv2d( 118 | int(256 * alpha), 119 | int(512 * alpha), 120 | 3, 121 | stride=2, 122 | padding=1, 123 | bias=False 124 | ), 125 | 126 | DepthSeperabelConv2d( 127 | int(512 * alpha), 128 | int(512 * alpha), 129 | 3, 130 | padding=1, 131 | bias=False 132 | ), 133 | DepthSeperabelConv2d( 134 | int(512 * alpha), 135 | int(512 * alpha), 136 | 3, 137 | padding=1, 138 | bias=False 139 | ), 140 | DepthSeperabelConv2d( 141 | int(512 * alpha), 142 | int(512 * alpha), 143 | 3, 144 | padding=1, 145 | bias=False 146 | ), 147 | DepthSeperabelConv2d( 148 | int(512 * alpha), 149 | int(512 * alpha), 150 | 3, 151 | padding=1, 152 | bias=False 153 | ), 154 | DepthSeperabelConv2d( 155 | int(512 * alpha), 156 | int(512 * alpha), 157 | 3, 158 | padding=1, 159 | bias=False 160 | ) 161 | ) 162 | 163 | #downsample 164 | self.conv4 = nn.Sequential( 165 | DepthSeperabelConv2d( 166 | int(512 * alpha), 167 | int(1024 * alpha), 168 | 3, 169 | stride=2, 170 | padding=1, 171 | bias=False 172 | ), 173 | DepthSeperabelConv2d( 174 | int(1024 * alpha), 175 | int(1024 * alpha), 176 | 3, 177 | padding=1, 178 | bias=False 179 | ) 180 | ) 181 | 182 | self.fc = nn.Linear(int(1024 * alpha), class_num) 183 | self.avg = nn.AdaptiveAvgPool2d(1) 184 | 185 | def forward(self, x): 186 | x = self.stem(x) 187 | 188 | x = self.conv1(x) 189 | x = self.conv2(x) 190 | x = self.conv3(x) 191 | x = self.conv4(x) 192 | 193 | x = self.avg(x) 194 | x = x.view(x.size(0), -1) 195 | x = self.fc(x) 196 | return x 197 | 198 | 199 | def mobilenet(class_num=100): 200 | alpha=1 201 | return MobileNet(alpha, class_num) 202 | 203 | 204 | if __name__ == "__main__": 205 | 206 | 207 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 208 | image_size = 256 209 | x = torch.Tensor(1, 1, image_size, image_size) 210 | x = x.to(device) 211 | print("x size: {}".format(x.size())) 212 | 213 | model = mobilenet(class_num=2).to(device) 214 | 215 | 216 | out1 = model(x) 217 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/two_d/nasnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class SeperableConv2d(nn.Module): 5 | 6 | def __init__(self, input_channels, output_channels, kernel_size, **kwargs): 7 | 8 | super().__init__() 9 | self.depthwise = nn.Conv2d( 10 | input_channels, 11 | input_channels, 12 | kernel_size, 13 | groups=input_channels, 14 | **kwargs 15 | ) 16 | 17 | self.pointwise = nn.Conv2d( 18 | input_channels, 19 | output_channels, 20 | 1 21 | ) 22 | def forward(self, x): 23 | x = self.depthwise(x) 24 | x = self.pointwise(x) 25 | 26 | return x 27 | 28 | class SeperableBranch(nn.Module): 29 | 30 | def __init__(self, input_channels, output_channels, kernel_size, **kwargs): 31 | """Adds 2 blocks of [relu-separable conv-batchnorm].""" 32 | super().__init__() 33 | self.block1 = nn.Sequential( 34 | nn.ReLU(), 35 | SeperableConv2d(input_channels, output_channels, kernel_size, **kwargs), 36 | nn.BatchNorm2d(output_channels) 37 | ) 38 | 39 | self.block2 = nn.Sequential( 40 | nn.ReLU(), 41 | SeperableConv2d(output_channels, output_channels, kernel_size, stride=1, padding=int(kernel_size / 2)), 42 | nn.BatchNorm2d(output_channels) 43 | ) 44 | 45 | def forward(self, x): 46 | x = self.block1(x) 47 | x = self.block2(x) 48 | 49 | return x 50 | 51 | class Fit(nn.Module): 52 | """Make the cell outputs compatible 53 | Args: 54 | prev_filters: filter number of tensor prev, needs to be modified 55 | filters: filter number of normal cell branch output filters 56 | """ 57 | 58 | def __init__(self, prev_filters, filters): 59 | super().__init__() 60 | self.relu = nn.ReLU() 61 | 62 | self.p1 = nn.Sequential( 63 | nn.AvgPool2d(1, stride=2), 64 | nn.Conv2d(prev_filters, int(filters / 2), 1) 65 | ) 66 | 67 | #make sure there is no information loss 68 | self.p2 = nn.Sequential( 69 | nn.ConstantPad2d((0, 1, 0, 1), 0), 70 | nn.ConstantPad2d((-1, 0, -1, 0), 0), #cropping 71 | nn.AvgPool2d(1, stride=2), 72 | nn.Conv2d(prev_filters, int(filters / 2), 1) 73 | ) 74 | 75 | self.bn = nn.BatchNorm2d(filters) 76 | 77 | self.dim_reduce = nn.Sequential( 78 | nn.ReLU(), 79 | nn.Conv2d(prev_filters, filters, 1), 80 | nn.BatchNorm2d(filters) 81 | ) 82 | 83 | self.filters = filters 84 | 85 | def forward(self, inputs): 86 | x, prev = inputs 87 | if prev is None: 88 | return x 89 | 90 | #image size does not match 91 | elif x.size(2) != prev.size(2): 92 | prev = self.relu(prev) 93 | p1 = self.p1(prev) 94 | p2 = self.p2(prev) 95 | prev = torch.cat([p1, p2], 1) 96 | prev = self.bn(prev) 97 | 98 | elif prev.size(1) != self.filters: 99 | prev = self.dim_reduce(prev) 100 | 101 | return prev 102 | 103 | 104 | class NormalCell(nn.Module): 105 | 106 | def __init__(self, x_in, prev_in, output_channels): 107 | super().__init__() 108 | 109 | self.dem_reduce = nn.Sequential( 110 | nn.ReLU(), 111 | nn.Conv2d(x_in, output_channels, 1, bias=False), 112 | nn.BatchNorm2d(output_channels) 113 | ) 114 | 115 | self.block1_left = SeperableBranch( 116 | output_channels, 117 | output_channels, 118 | kernel_size=3, 119 | padding=1, 120 | bias=False 121 | ) 122 | self.block1_right = nn.Sequential() 123 | 124 | self.block2_left = SeperableBranch( 125 | output_channels, 126 | output_channels, 127 | kernel_size=3, 128 | padding=1, 129 | bias=False 130 | ) 131 | self.block2_right = SeperableBranch( 132 | output_channels, 133 | output_channels, 134 | kernel_size=5, 135 | padding=2, 136 | bias=False 137 | ) 138 | 139 | self.block3_left = nn.AvgPool2d(3, stride=1, padding=1) 140 | self.block3_right = nn.Sequential() 141 | 142 | self.block4_left = nn.AvgPool2d(3, stride=1, padding=1) 143 | self.block4_right = nn.AvgPool2d(3, stride=1, padding=1) 144 | 145 | self.block5_left = SeperableBranch( 146 | output_channels, 147 | output_channels, 148 | kernel_size=5, 149 | padding=2, 150 | bias=False 151 | ) 152 | self.block5_right = SeperableBranch( 153 | output_channels, 154 | output_channels, 155 | kernel_size=3, 156 | padding=1, 157 | bias=False 158 | ) 159 | 160 | self.fit = Fit(prev_in, output_channels) 161 | 162 | def forward(self, x): 163 | x, prev = x 164 | 165 | #return transformed x as new x, and original x as prev 166 | #only prev tensor needs to be modified 167 | prev = self.fit((x, prev)) 168 | 169 | h = self.dem_reduce(x) 170 | 171 | x1 = self.block1_left(h) + self.block1_right(h) 172 | x2 = self.block2_left(prev) + self.block2_right(h) 173 | x3 = self.block3_left(h) + self.block3_right(h) 174 | x4 = self.block4_left(prev) + self.block4_right(prev) 175 | x5 = self.block5_left(prev) + self.block5_right(prev) 176 | 177 | return torch.cat([prev, x1, x2, x3, x4, x5], 1), x 178 | 179 | class ReductionCell(nn.Module): 180 | 181 | def __init__(self, x_in, prev_in, output_channels): 182 | super().__init__() 183 | 184 | self.dim_reduce = nn.Sequential( 185 | nn.ReLU(), 186 | nn.Conv2d(x_in, output_channels, 1), 187 | nn.BatchNorm2d(output_channels) 188 | ) 189 | 190 | #block1 191 | self.layer1block1_left = SeperableBranch(output_channels, output_channels, 7, stride=2, padding=3) 192 | self.layer1block1_right = SeperableBranch(output_channels, output_channels, 5, stride=2, padding=2) 193 | 194 | #block2 195 | self.layer1block2_left = nn.MaxPool2d(3, stride=2, padding=1) 196 | self.layer1block2_right = SeperableBranch(output_channels, output_channels, 7, stride=2, padding=3) 197 | 198 | #block3 199 | self.layer1block3_left = nn.AvgPool2d(3, 2, 1) 200 | self.layer1block3_right = SeperableBranch(output_channels, output_channels, 5, stride=2, padding=2) 201 | 202 | #block5 203 | self.layer2block1_left = nn.MaxPool2d(3, 2, 1) 204 | self.layer2block1_right = SeperableBranch(output_channels, output_channels, 3, stride=1, padding=1) 205 | 206 | #block4 207 | self.layer2block2_left = nn.AvgPool2d(3, 1, 1) 208 | self.layer2block2_right = nn.Sequential() 209 | 210 | self.fit = Fit(prev_in, output_channels) 211 | 212 | def forward(self, x): 213 | x, prev = x 214 | prev = self.fit((x, prev)) 215 | 216 | h = self.dim_reduce(x) 217 | 218 | layer1block1 = self.layer1block1_left(prev) + self.layer1block1_right(h) 219 | layer1block2 = self.layer1block2_left(h) + self.layer1block2_right(prev) 220 | layer1block3 = self.layer1block3_left(h) + self.layer1block3_right(prev) 221 | layer2block1 = self.layer2block1_left(h) + self.layer2block1_right(layer1block1) 222 | layer2block2 = self.layer2block2_left(layer1block1) + self.layer2block2_right(layer1block2) 223 | 224 | return torch.cat([ 225 | layer1block2, #https://github.com/keras-team/keras-applications/blob/master/keras_applications/nasnet.py line 739 226 | layer1block3, 227 | layer2block1, 228 | layer2block2 229 | ], 1), x 230 | 231 | 232 | class NasNetA(nn.Module): 233 | 234 | def __init__(self, repeat_cell_num, reduction_num, filters, stemfilter, class_num=100): 235 | super().__init__() 236 | 237 | self.stem = nn.Sequential( 238 | nn.Conv2d(1, stemfilter, 3, padding=1, bias=False), 239 | nn.BatchNorm2d(stemfilter) 240 | ) 241 | 242 | self.prev_filters = stemfilter 243 | self.x_filters = stemfilter 244 | self.filters = filters 245 | 246 | self.cell_layers = self._make_layers(repeat_cell_num, reduction_num) 247 | 248 | self.relu = nn.ReLU() 249 | self.avg = nn.AdaptiveAvgPool2d(1) 250 | self.fc = nn.Linear(self.filters * 6, class_num) 251 | 252 | 253 | def _make_normal(self, block, repeat, output): 254 | """make normal cell 255 | Args: 256 | block: cell type 257 | repeat: number of repeated normal cell 258 | output: output filters for each branch in normal cell 259 | Returns: 260 | stacked normal cells 261 | """ 262 | 263 | layers = [] 264 | for r in range(repeat): 265 | layers.append(block(self.x_filters, self.prev_filters, output)) 266 | self.prev_filters = self.x_filters 267 | self.x_filters = output * 6 #concatenate 6 branches 268 | 269 | return layers 270 | 271 | def _make_reduction(self, block, output): 272 | """make normal cell 273 | Args: 274 | block: cell type 275 | output: output filters for each branch in reduction cell 276 | Returns: 277 | reduction cell 278 | """ 279 | 280 | reduction = block(self.x_filters, self.prev_filters, output) 281 | self.prev_filters = self.x_filters 282 | self.x_filters = output * 4 #stack for 4 branches 283 | 284 | return reduction 285 | 286 | def _make_layers(self, repeat_cell_num, reduction_num): 287 | 288 | layers = [] 289 | for i in range(reduction_num): 290 | 291 | layers.extend(self._make_normal(NormalCell, repeat_cell_num, self.filters)) 292 | self.filters *= 2 293 | layers.append(self._make_reduction(ReductionCell, self.filters)) 294 | 295 | layers.extend(self._make_normal(NormalCell, repeat_cell_num, self.filters)) 296 | 297 | return nn.Sequential(*layers) 298 | 299 | 300 | def forward(self, x): 301 | 302 | x = self.stem(x) 303 | prev = None 304 | x, prev = self.cell_layers((x, prev)) 305 | x = self.relu(x) 306 | x = self.avg(x) 307 | x = x.view(x.size(0), -1) 308 | x = self.fc(x) 309 | 310 | return x 311 | 312 | 313 | def nasnet(class_num): 314 | 315 | #stem filters must be 44, it's a pytorch workaround, cant change to other number 316 | return NasNetA(4, 2, 44, 44,class_num=class_num) 317 | 318 | 319 | if __name__ == "__main__": 320 | 321 | 322 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 323 | image_size = 256 324 | x = torch.Tensor(1, 1, image_size, image_size) 325 | x = x.to(device) 326 | print("x size: {}".format(x.size())) 327 | 328 | model = nasnet(class_num=2).to(device) 329 | 330 | 331 | out1 = model(x) 332 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/two_d/resnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class BasicBlock(nn.Module): 5 | """Basic Block for resnet 18 and resnet 34 6 | """ 7 | 8 | #BasicBlock and BottleNeck block 9 | #have different output size 10 | #we use class attribute expansion 11 | #to distinct 12 | expansion = 1 13 | 14 | def __init__(self, in_channels, out_channels, stride=1): 15 | super().__init__() 16 | 17 | #residual function 18 | self.residual_function = nn.Sequential( 19 | nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False), 20 | nn.BatchNorm2d(out_channels), 21 | nn.ReLU(inplace=True), 22 | nn.Conv2d(out_channels, out_channels * BasicBlock.expansion, kernel_size=3, padding=1, bias=False), 23 | nn.BatchNorm2d(out_channels * BasicBlock.expansion) 24 | ) 25 | 26 | #shortcut 27 | self.shortcut = nn.Sequential() 28 | 29 | #the shortcut output dimension is not the same with residual function 30 | #use 1*1 convolution to match the dimension 31 | if stride != 1 or in_channels != BasicBlock.expansion * out_channels: 32 | self.shortcut = nn.Sequential( 33 | nn.Conv2d(in_channels, out_channels * BasicBlock.expansion, kernel_size=1, stride=stride, bias=False), 34 | nn.BatchNorm2d(out_channels * BasicBlock.expansion) 35 | ) 36 | 37 | def forward(self, x): 38 | return nn.ReLU(inplace=True)(self.residual_function(x) + self.shortcut(x)) 39 | 40 | class BottleNeck(nn.Module): 41 | """Residual block for resnet over 50 layers 42 | """ 43 | expansion = 4 44 | def __init__(self, in_channels, out_channels, stride=1): 45 | super().__init__() 46 | self.residual_function = nn.Sequential( 47 | nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), 48 | nn.BatchNorm2d(out_channels), 49 | nn.ReLU(inplace=True), 50 | nn.Conv2d(out_channels, out_channels, stride=stride, kernel_size=3, padding=1, bias=False), 51 | nn.BatchNorm2d(out_channels), 52 | nn.ReLU(inplace=True), 53 | nn.Conv2d(out_channels, out_channels * BottleNeck.expansion, kernel_size=1, bias=False), 54 | nn.BatchNorm2d(out_channels * BottleNeck.expansion), 55 | ) 56 | 57 | self.shortcut = nn.Sequential() 58 | 59 | if stride != 1 or in_channels != out_channels * BottleNeck.expansion: 60 | self.shortcut = nn.Sequential( 61 | nn.Conv2d(in_channels, out_channels * BottleNeck.expansion, stride=stride, kernel_size=1, bias=False), 62 | nn.BatchNorm2d(out_channels * BottleNeck.expansion) 63 | ) 64 | 65 | def forward(self, x): 66 | return nn.ReLU(inplace=True)(self.residual_function(x) + self.shortcut(x)) 67 | 68 | class ResNet(nn.Module): 69 | 70 | def __init__(self, block, num_block, num_classes=100): 71 | super().__init__() 72 | 73 | self.in_channels = 64 74 | 75 | self.conv1 = nn.Sequential( 76 | nn.Conv2d(1, 64, kernel_size=3, padding=1, bias=False), 77 | nn.BatchNorm2d(64), 78 | nn.ReLU(inplace=True)) 79 | #we use a different inputsize than the original paper 80 | #so conv2_x's stride is 1 81 | self.conv2_x = self._make_layer(block, 64, num_block[0], 1) 82 | self.conv3_x = self._make_layer(block, 128, num_block[1], 2) 83 | self.conv4_x = self._make_layer(block, 256, num_block[2], 2) 84 | self.conv5_x = self._make_layer(block, 512, num_block[3], 2) 85 | self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) 86 | self.fc = nn.Linear(512 * block.expansion, num_classes) 87 | 88 | def _make_layer(self, block, out_channels, num_blocks, stride): 89 | """make resnet layers(by layer i didnt mean this 'layer' was the 90 | same as a neuron netowork layer, ex. conv layer), one layer may 91 | contain more than one residual block 92 | Args: 93 | block: block type, basic block or bottle neck block 94 | out_channels: output depth channel number of this layer 95 | num_blocks: how many blocks per layer 96 | stride: the stride of the first block of this layer 97 | Return: 98 | return a resnet layer 99 | """ 100 | 101 | # we have num_block blocks per layer, the first block 102 | # could be 1 or 2, other blocks would always be 1 103 | strides = [stride] + [1] * (num_blocks - 1) 104 | layers = [] 105 | for stride in strides: 106 | layers.append(block(self.in_channels, out_channels, stride)) 107 | self.in_channels = out_channels * block.expansion 108 | 109 | return nn.Sequential(*layers) 110 | 111 | def forward(self, x): 112 | output = self.conv1(x) 113 | output = self.conv2_x(output) 114 | output = self.conv3_x(output) 115 | output = self.conv4_x(output) 116 | output = self.conv5_x(output) 117 | output = self.avg_pool(output) 118 | output = output.view(output.size(0), -1) 119 | output = self.fc(output) 120 | 121 | return output 122 | 123 | def resnet18(num_classes): 124 | """ return a ResNet 18 object 125 | """ 126 | return ResNet(BasicBlock, [2, 2, 2, 2],num_classes=num_classes) 127 | 128 | def resnet34(num_classes): 129 | """ return a ResNet 34 object 130 | """ 131 | return ResNet(BasicBlock, [3, 4, 6, 3],num_classes=num_classes) 132 | 133 | def resnet50(num_classes): 134 | """ return a ResNet 50 object 135 | """ 136 | return ResNet(BottleNeck, [3, 4, 6, 3],num_classes=num_classes) 137 | 138 | def resnet101(num_classes): 139 | """ return a ResNet 101 object 140 | """ 141 | return ResNet(BottleNeck, [3, 4, 23, 3],num_classes=num_classes) 142 | 143 | def resnet152(num_classes): 144 | """ return a ResNet 152 object 145 | """ 146 | return ResNet(BottleNeck, [3, 8, 36, 3],num_classes=num_classes) 147 | 148 | 149 | 150 | if __name__ == "__main__": 151 | 152 | 153 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 154 | image_size = 256 155 | x = torch.Tensor(1, 1, image_size, image_size) 156 | x = x.to(device) 157 | print("x size: {}".format(x.size())) 158 | 159 | model = resnet152(num_classes=2).to(device) 160 | 161 | 162 | out1 = model(x) 163 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/two_d/resnext.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | #only implements ResNext bottleneck c 7 | 8 | 9 | #"""This strategy exposes a new dimension, which we call “cardinality” 10 | #(the size of the set of transformations), as an essential factor 11 | #in addition to the dimensions of depth and width.""" 12 | CARDINALITY = 32 13 | DEPTH = 4 14 | BASEWIDTH = 64 15 | 16 | #"""The grouped convolutional layer in Fig. 3(c) performs 32 groups 17 | #of convolutions whose input and output channels are 4-dimensional. 18 | #The grouped convolutional layer concatenates them as the outputs 19 | #of the layer.""" 20 | 21 | class ResNextBottleNeckC(nn.Module): 22 | 23 | def __init__(self, in_channels, out_channels, stride): 24 | super().__init__() 25 | 26 | C = CARDINALITY #How many groups a feature map was splitted into 27 | 28 | #"""We note that the input/output width of the template is fixed as 29 | #256-d (Fig. 3), We note that the input/output width of the template 30 | #is fixed as 256-d (Fig. 3), and all widths are dou- bled each time 31 | #when the feature map is subsampled (see Table 1).""" 32 | D = int(DEPTH * out_channels / BASEWIDTH) #number of channels per group 33 | self.split_transforms = nn.Sequential( 34 | nn.Conv2d(in_channels, C * D, kernel_size=1, groups=C, bias=False), 35 | nn.BatchNorm2d(C * D), 36 | nn.ReLU(inplace=True), 37 | nn.Conv2d(C * D, C * D, kernel_size=3, stride=stride, groups=C, padding=1, bias=False), 38 | nn.BatchNorm2d(C * D), 39 | nn.ReLU(inplace=True), 40 | nn.Conv2d(C * D, out_channels * 4, kernel_size=1, bias=False), 41 | nn.BatchNorm2d(out_channels * 4), 42 | ) 43 | 44 | self.shortcut = nn.Sequential() 45 | 46 | if stride != 1 or in_channels != out_channels * 4: 47 | self.shortcut = nn.Sequential( 48 | nn.Conv2d(in_channels, out_channels * 4, stride=stride, kernel_size=1, bias=False), 49 | nn.BatchNorm2d(out_channels * 4) 50 | ) 51 | 52 | def forward(self, x): 53 | return F.relu(self.split_transforms(x) + self.shortcut(x)) 54 | 55 | class ResNext(nn.Module): 56 | 57 | def __init__(self, block, num_blocks, class_names=100): 58 | super().__init__() 59 | self.in_channels = 64 60 | 61 | self.conv1 = nn.Sequential( 62 | nn.Conv2d(1, 64, 3, stride=1, padding=1, bias=False), 63 | nn.BatchNorm2d(64), 64 | nn.ReLU(inplace=True) 65 | ) 66 | 67 | self.conv2 = self._make_layer(block, num_blocks[0], 64, 1) 68 | self.conv3 = self._make_layer(block, num_blocks[1], 128, 2) 69 | self.conv4 = self._make_layer(block, num_blocks[2], 256, 2) 70 | self.conv5 = self._make_layer(block, num_blocks[3], 512, 2) 71 | self.avg = nn.AdaptiveAvgPool2d((1, 1)) 72 | self.fc = nn.Linear(512 * 4, class_names) 73 | 74 | def forward(self, x): 75 | x = self.conv1(x) 76 | x = self.conv2(x) 77 | x = self.conv3(x) 78 | x = self.conv4(x) 79 | x = self.conv5(x) 80 | x = self.avg(x) 81 | x = x.view(x.size(0), -1) 82 | x = self.fc(x) 83 | return x 84 | 85 | def _make_layer(self, block, num_block, out_channels, stride): 86 | """Building resnext block 87 | Args: 88 | block: block type(default resnext bottleneck c) 89 | num_block: number of blocks per layer 90 | out_channels: output channels per block 91 | stride: block stride 92 | Returns: 93 | a resnext layer 94 | """ 95 | strides = [stride] + [1] * (num_block - 1) 96 | layers = [] 97 | for stride in strides: 98 | layers.append(block(self.in_channels, out_channels, stride)) 99 | self.in_channels = out_channels * 4 100 | 101 | return nn.Sequential(*layers) 102 | 103 | def resnext50(class_names): 104 | """ return a resnext50(c32x4d) network 105 | """ 106 | return ResNext(ResNextBottleNeckC, [3, 4, 6, 3], class_names=class_names) 107 | 108 | def resnext101(class_names): 109 | """ return a resnext101(c32x4d) network 110 | """ 111 | return ResNext(ResNextBottleNeckC, [3, 4, 23, 3], class_names=class_names) 112 | 113 | def resnext152(class_names): 114 | """ return a resnext101(c32x4d) network 115 | """ 116 | return ResNext(ResNextBottleNeckC, [3, 4, 36, 3], class_names=class_names) 117 | 118 | 119 | if __name__ == "__main__": 120 | 121 | 122 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 123 | image_size = 256 124 | x = torch.Tensor(1, 1, image_size, image_size) 125 | x = x.to(device) 126 | print("x size: {}".format(x.size())) 127 | 128 | model = resnext152(class_names=2).to(device) 129 | 130 | 131 | out1 = model(x) 132 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /models/two_d/vggnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | cfg = { 5 | 'A' : [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 6 | 'B' : [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 7 | 'D' : [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 8 | 'E' : [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'] 9 | } 10 | 11 | class VGG(nn.Module): 12 | 13 | def __init__(self, features, num_class=100): 14 | super().__init__() 15 | self.features = features 16 | 17 | self.classifier = nn.Sequential( 18 | nn.Linear(512*7*7, 4096), 19 | nn.ReLU(inplace=True), 20 | nn.Dropout(), 21 | nn.Linear(4096, 4096), 22 | nn.ReLU(inplace=True), 23 | nn.Dropout(), 24 | nn.Linear(4096, num_class) 25 | ) 26 | 27 | def forward(self, x): 28 | output = self.features(x) 29 | output = output.view(output.size()[0], -1) 30 | output = self.classifier(output) 31 | 32 | return output 33 | 34 | def make_layers(cfg, batch_norm=False): 35 | layers = [] 36 | 37 | input_channel = 1 38 | for l in cfg: 39 | if l == 'M': 40 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 41 | continue 42 | 43 | layers += [nn.Conv2d(input_channel, l, kernel_size=3, padding=1)] 44 | 45 | if batch_norm: 46 | layers += [nn.BatchNorm2d(l)] 47 | 48 | layers += [nn.ReLU(inplace=True)] 49 | input_channel = l 50 | 51 | return nn.Sequential(*layers) 52 | 53 | def vgg11_bn(num_class): 54 | return VGG(make_layers(cfg['A'], batch_norm=True),num_class=num_class) 55 | 56 | def vgg13_bn(num_class): 57 | return VGG(make_layers(cfg['B'], batch_norm=True),num_class=num_class) 58 | 59 | def vgg16_bn(num_class): 60 | return VGG(make_layers(cfg['D'], batch_norm=True),num_class=num_class) 61 | 62 | def vgg19_bn(num_class): 63 | return VGG(make_layers(cfg['E'], batch_norm=True),num_class=num_class) 64 | 65 | 66 | if __name__ == "__main__": 67 | 68 | 69 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 70 | image_size = 224 71 | x = torch.Tensor(1, 1, image_size, image_size) 72 | x = x.to(device) 73 | print("x size: {}".format(x.size())) 74 | 75 | model = vgg16_bn(num_class=2).to(device) 76 | 77 | 78 | out1 = model(x) 79 | print("out size: {}".format(out1.size())) -------------------------------------------------------------------------------- /utils/convolution.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | 4 | 5 | PADDING_MODES = { 6 | 'reflect': 'Reflection', 7 | 'replicate': 'Replication', 8 | 'constant': 'Zero', 9 | } 10 | 11 | 12 | class ConvolutionalBlock(nn.Module): 13 | def __init__( 14 | self, 15 | in_channels, 16 | out_channels, 17 | dilation, 18 | dimensions, 19 | batch_norm=True, 20 | instance_norm=False, 21 | norm_affine=True, 22 | padding_mode='constant', 23 | preactivation=True, 24 | kernel_size=3, 25 | activation=True, 26 | ): 27 | assert padding_mode in PADDING_MODES.keys() 28 | assert not (batch_norm and instance_norm) 29 | super().__init__() 30 | 31 | if dimensions == 2: 32 | # pylint: disable=not-callable 33 | class_name = '{}Pad2d'.format(PADDING_MODES[padding_mode]) 34 | padding_class = getattr(nn, class_name) 35 | padding_instance = padding_class(dilation) 36 | elif dimensions == 3: 37 | padding_instance = Pad3d(dilation, padding_mode) 38 | conv_class = nn.Conv2d if dimensions == 2 else nn.Conv3d 39 | 40 | if batch_norm: 41 | norm_class = nn.BatchNorm2d if dimensions == 2 else nn.BatchNorm3d 42 | if instance_norm: 43 | norm_class = nn.InstanceNorm2d if dimensions == 2 else nn.InstanceNorm3d 44 | 45 | layers = nn.ModuleList() 46 | 47 | if preactivation: 48 | if batch_norm or instance_norm: 49 | layers.append(norm_class(in_channels, affine=norm_affine)) 50 | if activation: 51 | layers.append(nn.ReLU()) 52 | 53 | if kernel_size > 1: 54 | layers.append(padding_instance) 55 | 56 | use_bias = not (instance_norm or batch_norm) 57 | conv_layer = conv_class( 58 | in_channels, 59 | out_channels, 60 | kernel_size=kernel_size, 61 | dilation=dilation, 62 | bias=use_bias, 63 | ) 64 | layers.append(conv_layer) 65 | 66 | if not preactivation: 67 | if batch_norm or instance_norm: 68 | layers.append(norm_class(out_channels, affine=norm_affine)) 69 | if activation: 70 | layers.append(nn.ReLU()) 71 | 72 | self.convolutional_block = nn.Sequential(*layers) 73 | 74 | def forward(self, x): 75 | return self.convolutional_block(x) 76 | 77 | 78 | class Pad3d(nn.Module): 79 | def __init__(self, pad, mode): 80 | assert mode in PADDING_MODES.keys() 81 | super().__init__() 82 | self.pad = 6 * [pad] 83 | self.mode = mode 84 | 85 | def forward(self, x): 86 | return F.pad(x, self.pad, self.mode) -------------------------------------------------------------------------------- /utils/dilation.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from residual import ResidualBlock 3 | 4 | 5 | class DilationBlock(nn.Module): 6 | def __init__( 7 | self, 8 | in_channels, 9 | out_channels, 10 | dilation, 11 | dimensions, 12 | layers_per_block=2, 13 | num_residual_blocks=3, 14 | batch_norm=True, 15 | instance_norm=False, 16 | residual=True, 17 | padding_mode='constant', 18 | ): 19 | super().__init__() 20 | self.in_channels = in_channels 21 | self.out_channels = out_channels 22 | residual_blocks = nn.ModuleList() 23 | for _ in range(num_residual_blocks): 24 | residual_block = ResidualBlock( 25 | in_channels, 26 | out_channels, 27 | layers_per_block, 28 | dilation, 29 | dimensions, 30 | batch_norm=batch_norm, 31 | instance_norm=instance_norm, 32 | residual=residual, 33 | padding_mode=padding_mode, 34 | ) 35 | residual_blocks.append(residual_block) 36 | in_channels = out_channels 37 | self.dilation_block = nn.Sequential(*residual_blocks) 38 | 39 | def forward(self, x): 40 | return self.dilation_block(x) -------------------------------------------------------------------------------- /utils/metric.py: -------------------------------------------------------------------------------- 1 | import torchio as tio 2 | from pathlib import Path 3 | import torch 4 | import numpy as np 5 | import copy 6 | 7 | def metric(gt,pred): 8 | preds = pred.detach().numpy() 9 | gts = gt.detach().numpy() 10 | 11 | 12 | 13 | pred = preds.astype(int) # float data does not support bit_and and bit_or 14 | gdth = gts.astype(int) # float data does not support bit_and and bit_or 15 | fp_array = copy.deepcopy(pred) # keep pred unchanged 16 | fn_array = copy.deepcopy(gdth) 17 | gdth_sum = np.sum(gdth) 18 | pred_sum = np.sum(pred) 19 | intersection = gdth & pred 20 | union = gdth | pred 21 | intersection_sum = np.count_nonzero(intersection) 22 | union_sum = np.count_nonzero(union) 23 | 24 | tp_array = intersection 25 | 26 | tmp = pred - gdth 27 | fp_array[tmp < 1] = 0 28 | 29 | tmp2 = gdth - pred 30 | fn_array[tmp2 < 1] = 0 31 | 32 | tn_array = np.ones(gdth.shape) - union 33 | 34 | tp, fp, fn, tn = np.sum(tp_array), np.sum(fp_array), np.sum(fn_array), np.sum(tn_array) 35 | 36 | smooth = 0.001 37 | precision = tp / (pred_sum + smooth) 38 | recall = tp / (gdth_sum + smooth) 39 | 40 | false_positive_rate = fp / (fp + tn + smooth) 41 | false_negtive_rate = fn / (fn + tp + smooth) 42 | 43 | jaccard = intersection_sum / (union_sum + smooth) 44 | dice = 2 * intersection_sum / (gdth_sum + pred_sum + smooth) 45 | 46 | return false_positive_rate,false_negtive_rate,dice 47 | -------------------------------------------------------------------------------- /utils/residual.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from convolution import ConvolutionalBlock 5 | 6 | 7 | BATCH_DIM = 0 8 | CHANNELS_DIM = 1 9 | 10 | 11 | class ResidualBlock(nn.Module): 12 | def __init__( 13 | self, 14 | in_channels, 15 | out_channels, 16 | num_layers, 17 | dilation, 18 | dimensions, 19 | batch_norm=True, 20 | instance_norm=False, 21 | residual=True, 22 | residual_type='pad', 23 | padding_mode='constant', 24 | ): 25 | assert residual_type in ('pad', 'project') 26 | super().__init__() 27 | self.residual = residual 28 | self.change_dimension = in_channels != out_channels 29 | self.residual_type = residual_type 30 | self.dimensions = dimensions 31 | if self.change_dimension: 32 | if residual_type == 'project': 33 | conv_class = nn.Conv2d if dimensions == 2 else nn.Conv3d 34 | self.change_dim_layer = conv_class( 35 | in_channels, 36 | out_channels, 37 | kernel_size=1, 38 | dilation=dilation, 39 | bias=False, # as in NiftyNet and PyTorch's ResNet model 40 | ) 41 | 42 | conv_blocks = nn.ModuleList() 43 | for _ in range(num_layers): 44 | conv_block = ConvolutionalBlock( 45 | in_channels, 46 | out_channels, 47 | dilation, 48 | dimensions, 49 | batch_norm=batch_norm, 50 | instance_norm=instance_norm, 51 | padding_mode=padding_mode, 52 | ) 53 | conv_blocks.append(conv_block) 54 | in_channels = out_channels 55 | self.residual_block = nn.Sequential(*conv_blocks) 56 | 57 | def forward(self, x): 58 | """ 59 | From the original ResNet paper, page 4: 60 | "When the dimensions increase, we consider two options: 61 | (A) The shortcut still performs identity mapping, 62 | with extra zero entries padded for increasing dimensions. 63 | This option introduces no extra parameter 64 | (B) The projection shortcut in Eqn.(2) is used to 65 | match dimensions (done by 1x1 convolutions). 66 | For both options, when the shortcuts go across feature maps of 67 | two sizes, they are performed with a stride of 2." 68 | """ 69 | out = self.residual_block(x) 70 | if self.residual: 71 | if self.change_dimension: 72 | if self.residual_type == 'project': 73 | x = self.change_dim_layer(x) 74 | elif self.residual_type == 'pad': 75 | batch_size = x.shape[BATCH_DIM] 76 | x_channels = x.shape[CHANNELS_DIM] 77 | out_channels = out.shape[CHANNELS_DIM] 78 | spatial_dims = x.shape[2:] 79 | diff_channels = out_channels - x_channels 80 | zeros_half = x.new_zeros( 81 | batch_size, diff_channels // 2, *spatial_dims) 82 | x = torch.cat((zeros_half, x, zeros_half), 83 | dim=CHANNELS_DIM) 84 | out = x + out 85 | return out --------------------------------------------------------------------------------