Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,450 Bytes
cfb7702 |
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 |
import pytorch_lightning as pl
import models
from systems.utils import parse_optimizer, parse_scheduler, update_module_step
from utils.mixins import SaverMixin
from utils.misc import config_to_primitive, get_rank
class BaseSystem(pl.LightningModule, SaverMixin):
"""
Two ways to print to console:
1. self.print: correctly handle progress bar
2. rank_zero_info: use the logging module
"""
def __init__(self, config):
super().__init__()
self.config = config
self.rank = get_rank()
self.prepare()
self.model = models.make(self.config.model.name, self.config.model)
def prepare(self):
pass
def forward(self, batch):
raise NotImplementedError
def C(self, value):
if isinstance(value, int) or isinstance(value, float):
pass
else:
value = config_to_primitive(value)
if not isinstance(value, list):
raise TypeError('Scalar specification only supports list, got', type(value))
if len(value) == 3:
value = [0] + value
assert len(value) == 4
start_step, start_value, end_value, end_step = value
if isinstance(end_step, int):
current_step = self.global_step
value = start_value + (end_value - start_value) * max(min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0)
elif isinstance(end_step, float):
current_step = self.current_epoch
value = start_value + (end_value - start_value) * max(min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0)
return value
def preprocess_data(self, batch, stage):
pass
"""
Implementing on_after_batch_transfer of DataModule does the same.
But on_after_batch_transfer does not support DP.
"""
def on_train_batch_start(self, batch, batch_idx, unused=0):
self.dataset = self.trainer.datamodule.train_dataloader().dataset
self.preprocess_data(batch, 'train')
update_module_step(self.model, self.current_epoch, self.global_step)
def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
self.dataset = self.trainer.datamodule.val_dataloader().dataset
self.preprocess_data(batch, 'validation')
update_module_step(self.model, self.current_epoch, self.global_step)
def on_test_batch_start(self, batch, batch_idx, dataloader_idx):
self.dataset = self.trainer.datamodule.test_dataloader().dataset
self.preprocess_data(batch, 'test')
update_module_step(self.model, self.current_epoch, self.global_step)
def on_predict_batch_start(self, batch, batch_idx, dataloader_idx):
self.dataset = self.trainer.datamodule.predict_dataloader().dataset
self.preprocess_data(batch, 'predict')
update_module_step(self.model, self.current_epoch, self.global_step)
def training_step(self, batch, batch_idx):
raise NotImplementedError
"""
# aggregate outputs from different devices (DP)
def training_step_end(self, out):
pass
"""
"""
# aggregate outputs from different iterations
def training_epoch_end(self, out):
pass
"""
def validation_step(self, batch, batch_idx):
raise NotImplementedError
"""
# aggregate outputs from different devices when using DP
def validation_step_end(self, out):
pass
"""
def validation_epoch_end(self, out):
"""
Gather metrics from all devices, compute mean.
Purge repeated results using data index.
"""
raise NotImplementedError
def test_step(self, batch, batch_idx):
raise NotImplementedError
def test_epoch_end(self, out):
"""
Gather metrics from all devices, compute mean.
Purge repeated results using data index.
"""
raise NotImplementedError
def export(self):
raise NotImplementedError
def configure_optimizers(self):
optim = parse_optimizer(self.config.system.optimizer, self.model)
ret = {
'optimizer': optim,
}
if 'scheduler' in self.config.system:
ret.update({
'lr_scheduler': parse_scheduler(self.config.system.scheduler, optim),
})
return ret
|