This repository has been archived by the owner on Mar 21, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 142
/
ssl_container.py
272 lines (246 loc) · 16.7 KB
/
ssl_container.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
# ------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
# ------------------------------------------------------------------------------------------
import logging
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union
import param
from pytorch_lightning import LightningModule
from yacs.config import CfgNode
from InnerEye.ML.SSL.datamodules_and_datasets.cifar_datasets import InnerEyeCIFAR10, InnerEyeCIFAR100
from InnerEye.ML.SSL.datamodules_and_datasets.cxr_datasets import CheXpert, CovidDataset, NIHCXR, RSNAKaggleCXR
from InnerEye.ML.SSL.datamodules_and_datasets.datamodules import CombinedDataModule, InnerEyeVisionDataModule
from InnerEye.ML.SSL.datamodules_and_datasets.transforms_utils import InnerEyeCIFARLinearHeadTransform, \
InnerEyeCIFARTrainTransform, \
get_ssl_transforms_from_config
from InnerEye.ML.SSL.encoders import get_encoder_output_dim
from InnerEye.ML.SSL.lightning_modules.byol.byol_module import BYOLInnerEye
from InnerEye.ML.SSL.lightning_modules.simclr_module import SimCLRInnerEye
from InnerEye.ML.SSL.lightning_modules.ssl_online_evaluator import SSLOnlineEvaluatorInnerEye
from InnerEye.ML.SSL.utils import SSLDataModuleType, SSLTrainingType, load_yaml_augmentation_config
from InnerEye.ML.lightning_container import LightningContainer
@dataclass
class DataModuleArgs:
augmentation_params: Optional[CfgNode]
dataset_name: str
dataset_path: Optional[Path]
batch_size: int
class EncoderName(Enum):
resnet18 = "resnet18"
resnet50 = "resnet50"
resnet101 = "resnet101"
densenet121 = "densenet121"
class SSLDatasetName(Enum):
CIFAR10 = "CIFAR10"
CIFAR100 = "CIFAR100"
RSNAKaggleCXR = "RSNAKaggleCXR"
NIHCXR = "NIHCXR"
CheXpert = "CheXpert"
Covid = "CovidDataset"
InnerEyeDataModuleTypes = Union[InnerEyeVisionDataModule, CombinedDataModule]
class SSLContainer(LightningContainer):
"""
This container is the based module to train an SSL model (either using BYOL or SimCLR).
To have an overview of the parameters available for configuring this container please check out the documentation
at docs/self_supervised_models.md.
Note that this container is also used as the base class for SSLImageClassifier (finetuning container) as they share
setup and datamodule methods.
"""
_SSLDataClassMappings = {SSLDatasetName.CIFAR10.value: InnerEyeCIFAR10,
SSLDatasetName.CIFAR100.value: InnerEyeCIFAR100,
SSLDatasetName.RSNAKaggleCXR.value: RSNAKaggleCXR,
SSLDatasetName.NIHCXR.value: NIHCXR,
SSLDatasetName.CheXpert.value: CheXpert,
SSLDatasetName.Covid.value: CovidDataset}
ssl_augmentation_config = param.ClassSelector(class_=Path, allow_None=True,
doc="The path to the yaml config defining the parameters of the "
"augmentations. Ignored for CIFAR10 example")
ssl_training_dataset_name = param.ClassSelector(class_=SSLDatasetName, doc="The name of the dataset")
ssl_training_batch_size = param.Integer(
doc="Total training batch size, will be divided across the number of gpus used for training. For example: if "
"you specify ssl_training_batch_size=1600 and use 4 nodes with 4 gpus each (i.e. total of 16 GPUs), "
"the code will provide a per-gpu batch size of 100")
ssl_training_type = param.ClassSelector(class_=SSLTrainingType, doc="Which algorithm to use for SSL training")
ssl_encoder = param.ClassSelector(class_=EncoderName, doc="Which encoder to use for SSL")
use_balanced_binary_loss_for_linear_head = param.Boolean(default=False,
doc="Whether to use a balanced loss for the training of "
"the linear head")
num_workers = param.Integer(default=6, doc="Number of workers to use for dataloader processes.")
is_debug_model = param.Boolean(default=False,
doc="If True, the training will be restricted to 1 batch per epoch."
"Used for debugging and tests.")
linear_head_augmentation_config = param.ClassSelector(class_=Path,
doc="The path to the yaml config for the linear head "
"augmentations")
linear_head_dataset_name = param.ClassSelector(class_=SSLDatasetName,
doc="Name of the dataset to use for the linear head training")
linear_head_batch_size = param.Integer(default=256, doc="Batch size for linear head tuning")
learning_rate_linear_head_during_ssl_training = param.Number(default=1e-4,
doc="Learning rate for linear head training during "
"SSL training.")
drop_last = param.Boolean(default=True, doc="If True drops the last incomplete batch")
def setup(self) -> None:
from InnerEye.ML.SSL.lightning_containers.ssl_image_classifier import SSLClassifierContainer
self.total_num_gpus = self.num_gpus_per_node * self.num_nodes
self._load_config()
# If you're using the same data for training and linear head, allow the user to specify the dataset only
# once. Or if you are doing just finetuning of linear head, the user should be able to specify dataset via
# azure_dataset_id/local_dataset instead of extra_dataset fields (as in this case we only use one dataset).
if ((self.linear_head_dataset_name == self.ssl_training_dataset_name) or isinstance(self,
SSLClassifierContainer)) \
and len(self.extra_local_dataset_paths) == 0 and self.local_dataset is not None:
self.extra_local_dataset_paths = [self.local_dataset]
self.datamodule_args = {SSLDataModuleType.LINEAR_HEAD:
DataModuleArgs(augmentation_params=self.classifier_augmentation_params,
dataset_name=self.linear_head_dataset_name.value,
dataset_path=self.extra_local_dataset_paths[0] if len(
self.extra_local_dataset_paths) > 0 else None,
batch_size=self.linear_head_batch_size)}
if self.ssl_training_dataset_name is not None:
self.datamodule_args.update(
{SSLDataModuleType.ENCODER: DataModuleArgs(augmentation_params=self.ssl_augmentation_params,
dataset_name=self.ssl_training_dataset_name.value,
dataset_path=self.local_dataset,
batch_size=self.ssl_training_batch_size)})
self.data_module: InnerEyeDataModuleTypes = self.get_data_module()
self.inference_on_val_set = False
self.inference_on_test_set = False
if self.perform_cross_validation:
raise NotImplementedError("Cross-validation logic is not implemented for this module.")
def _load_config(self) -> None:
# For Chest-XRay you need to specify the parameters of the augmentations via a config file.
self.ssl_augmentation_params = load_yaml_augmentation_config(
self.ssl_augmentation_config) if self.ssl_augmentation_config is not None \
else None
self.classifier_augmentation_params = load_yaml_augmentation_config(
self.linear_head_augmentation_config) if self.linear_head_augmentation_config is not None else \
self.ssl_augmentation_params
def create_model(self) -> LightningModule:
"""
This method must create the actual Lightning model that will be trained.
"""
# For small images like CIFAR, if using a resnet encoder, switch the first conv layer to a 3x3 kernel instead
# of a 7x7 conv layer.
use_7x7_first_conv_in_resnet = False if self.ssl_training_dataset_name.value.startswith("CIFAR") else True
if self.ssl_training_type == SSLTrainingType.SimCLR:
model: LightningModule = SimCLRInnerEye(encoder_name=self.ssl_encoder.value,
dataset_name=self.ssl_training_dataset_name.value,
use_7x7_first_conv_in_resnet=use_7x7_first_conv_in_resnet,
gpus=self.total_num_gpus,
num_samples=self.data_module.num_samples,
batch_size=self.data_module.batch_size,
lr=self.l_rate,
max_epochs=self.num_epochs)
elif self.ssl_training_type == SSLTrainingType.BYOL:
model = BYOLInnerEye(encoder_name=self.ssl_encoder.value,
num_samples=self.data_module.num_samples,
batch_size=self.data_module.batch_size,
learning_rate=self.l_rate,
use_7x7_first_conv_in_resnet=use_7x7_first_conv_in_resnet,
warmup_epochs=10,
max_epochs=self.num_epochs)
else:
raise ValueError(
f"Unknown value for ssl_training_type, should be {SSLTrainingType.SimCLR.value} or "
f"{SSLTrainingType.BYOL.value}. "
f"Found {self.ssl_training_type.value}")
model.hparams.update({'ssl_type': self.ssl_training_type.value,
"num_classes": self.data_module.num_classes})
self.encoder_output_dim = get_encoder_output_dim(model, self.data_module)
return model
def get_data_module(self) -> InnerEyeDataModuleTypes:
"""
Gets the data that is used for the training and validation steps.
Here we use different data loader for training of linear head and training of SSL model.
"""
if hasattr(self, "data_module"):
return self.data_module
encoder_data_module = self._create_ssl_data_modules(is_ssl_encoder_module=True)
linear_data_module = self._create_ssl_data_modules(is_ssl_encoder_module=False)
return CombinedDataModule(encoder_data_module, linear_data_module,
self.use_balanced_binary_loss_for_linear_head)
def _create_ssl_data_modules(self, is_ssl_encoder_module: bool) -> InnerEyeVisionDataModule:
"""
Returns torch lightning data module for encoder or linear head
:param is_ssl_encoder_module: whether to return the data module for SSL training or for linear head. If true,
:return transforms with two views per sample (batch like (img_v1, img_v2, label)). If False, return only one
view per sample but also return the index of the sample in the dataset (to make sure we don't use twice the same
batch in one training epoch (batch like (index, img_v1, label), as classifier dataloader expected to be shorter
than SSL training, hence CombinedDataloader might loop over data several times per epoch).
"""
datamodule_args = self.datamodule_args[SSLDataModuleType.ENCODER] if is_ssl_encoder_module else \
self.datamodule_args[SSLDataModuleType.LINEAR_HEAD]
train_transforms, val_transforms = self._get_transforms(datamodule_args.augmentation_params,
datamodule_args.dataset_name,
is_ssl_encoder_module)
batch_size_per_gpu = datamodule_args.batch_size // self.total_num_gpus if self.total_num_gpus > 0 else \
datamodule_args.batch_size
logging.info(f"Batch size per gpu: {batch_size_per_gpu}")
dm = InnerEyeVisionDataModule(dataset_cls=self._SSLDataClassMappings[datamodule_args.dataset_name],
return_index=not is_ssl_encoder_module, # index is only needed for linear head
train_transforms=train_transforms,
val_split=0.1,
val_transforms=val_transforms,
data_dir=str(datamodule_args.dataset_path),
batch_size=batch_size_per_gpu,
num_workers=self.num_workers,
seed=self.random_seed,
drop_last=self.drop_last)
dm.prepare_data()
dm.setup()
return dm
def _get_transforms(self, augmentation_config: Optional[CfgNode],
dataset_name: str,
is_ssl_encoder_module: bool) -> Tuple[Any, Any]:
"""
Returns the transformation pipeline for training and validation.
:param augmentation_config: optional yaml config defining strength of augmenentations. Ignored for CIFAR
examples.
:param dataset_name: name of the dataset, value has to be in SSLDatasetName, determines which transformation
pipeline to return.
:param is_ssl_encoder_module: if True the transformation pipeline will yield two versions of the image it is
applied on and it applies the same transformations for validation. If False, return only one transformation.
:return: training transformation pipeline and validation transformation pipeline.
"""
if dataset_name in [SSLDatasetName.RSNAKaggleCXR.value,
SSLDatasetName.NIHCXR.value,
SSLDatasetName.CheXpert.value,
SSLDatasetName.Covid.value]:
assert augmentation_config is not None
train_transforms, val_transforms = get_ssl_transforms_from_config(
augmentation_config,
return_two_views_per_sample=is_ssl_encoder_module,
use_training_augmentations_for_validation=is_ssl_encoder_module
)
elif dataset_name in [SSLDatasetName.CIFAR10.value, SSLDatasetName.CIFAR100.value]:
train_transforms = \
InnerEyeCIFARTrainTransform(32) if is_ssl_encoder_module else InnerEyeCIFARLinearHeadTransform(32)
val_transforms = \
InnerEyeCIFARTrainTransform(32) if is_ssl_encoder_module else InnerEyeCIFARLinearHeadTransform(32)
elif augmentation_config:
train_transforms, val_transforms = get_ssl_transforms_from_config(
augmentation_config,
return_two_views_per_sample=is_ssl_encoder_module,
use_training_augmentations_for_validation=is_ssl_encoder_module,
expand_channels=False,
)
logging.warning(f"Dataset {dataset_name} unknown. The config will be consumed by "
f"get_ssl_transforms() to create the augmentation pipeline, make sure "
f"the transformations in your configs are compatible. ")
else:
raise ValueError(f"Dataset {dataset_name} unknown and no config has been passed.")
return train_transforms, val_transforms
def get_trainer_arguments(self) -> Dict[str, Any]:
self.online_eval = SSLOnlineEvaluatorInnerEye(class_weights=self.data_module.class_weights, # type: ignore
z_dim=self.encoder_output_dim,
num_classes=self.data_module.num_classes, # type: ignore
dataset=self.linear_head_dataset_name.value, # type: ignore
drop_p=0.2,
learning_rate=self.learning_rate_linear_head_during_ssl_training)
trainer_kwargs: Dict[str, Any] = {"callbacks": self.online_eval}
if self.is_debug_model:
trainer_kwargs.update({"limit_train_batches": 1, "limit_val_batches": 1})
return trainer_kwargs