From b1858671410d73ce7c0e7b3a519497ca8983d28f Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 27 Oct 2021 10:14:06 +0800 Subject: [PATCH 01/30] added decorator to create a pytorch lightning model from torch --- .../vision/models/lightning_support.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 python/nano/src/bigdl/nano/pytorch/vision/models/lightning_support.py diff --git a/python/nano/src/bigdl/nano/pytorch/vision/models/lightning_support.py b/python/nano/src/bigdl/nano/pytorch/vision/models/lightning_support.py new file mode 100644 index 00000000000..72ded4158c1 --- /dev/null +++ b/python/nano/src/bigdl/nano/pytorch/vision/models/lightning_support.py @@ -0,0 +1,56 @@ +# +# Copyright 2016 The BigDL Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import pytorch_lightning as pl + + +def lightning_module(loss_creator, optimizer_creator, configs): + def lightning_from_torch(cls): + class LightningModel(pl.LightningModule): + def __init__(self, *args, **kwargs): + super().__init__() + self.model = cls(*args, **kwargs) + self.loss = loss_creator(configs) + self.optimizer = optimizer_creator(self.model, configs) + + def forward(self, batch): + # Handle different numbers of input for various models + nargs = self.model.forward.__code__.co_argcount + return self.model(*batch[:nargs - 1]) + + def training_step(self, batch, batch_idx): + x, y = batch + y_hat = self(batch) + loss = self.loss(y_hat, y) + return loss + + def validation_step(self, batch, batch_idx): + x, y = batch + y_hat = self(batch) + loss = self.loss(y_hat, y) + return loss + + def test_step(self, batch, batch_idx): + x, y = batch + y_hat = self(batch) + loss = self.loss(y_hat, y) + return loss + + def configure_optimizers(self): + return self.optimizer + + return LightningModel + + return lightning_from_torch From 8151c8a099acef843737a2077167fa1a88908f70 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 27 Oct 2021 14:04:02 +0800 Subject: [PATCH 02/30] added unit test for pytorch lightning decorator --- .../test/test_models_lightning_support.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 python/nano/test/test_models_lightning_support.py diff --git a/python/nano/test/test_models_lightning_support.py b/python/nano/test/test_models_lightning_support.py new file mode 100644 index 00000000000..3717b6421fd --- /dev/null +++ b/python/nano/test/test_models_lightning_support.py @@ -0,0 +1,57 @@ +# +# Copyright 2016 The BigDL Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +from unittest import TestCase + +import torch +from test._train_torch_lightning import train_torch_lightning +from torch import nn + +from bigdl.nano.pytorch.vision.models import lightning_support, vision + +config = { + 'lr': 0.01, + 'optim': 'Adam', +} +batch_size = 256 +num_workers = 0 +data_dir = os.path.join(os.path.dirname(__file__), "data") + + +def loss_creator(config): + return nn.CrossEntropyLoss() + + +def optimizer_creator(model, config): + return getattr(torch.optim, config.get("optim", "Adam"))(model.parameters(), lr=config.get("lr", 0.001)) + + +@lightning_support.lightning_module(loss_creator, optimizer_creator, config) +def resnet18(num_classes, pretrained=True, include_top=False, freeze=True): + backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) + output_size = backbone.get_output_size() + head = nn.Linear(output_size, num_classes) + return torch.nn.Sequential(backbone, head) + + +class TestModelsLightningSupport(TestCase): + num_classes = 10 + + def test_resnet18_ipex(self): + pl_resnet18 = resnet18(10, pretrained=True, include_top=False, freeze=True) + train_torch_lightning(pl_resnet18, batch_size, num_workers, data_dir, + use_orca_lite_trainer=True) From 5e4b992918ffa3657a7cc81a14ef1dcf4691823e Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 28 Oct 2021 15:30:15 +0800 Subject: [PATCH 03/30] refactoring - renaming, adding hints and docstring --- ...ning_support.py => lightning_extension.py} | 20 +++++++++++++++---- ....py => test_models_lightning_extension.py} | 5 +++-- 2 files changed, 19 insertions(+), 6 deletions(-) rename python/nano/src/bigdl/nano/pytorch/vision/models/{lightning_support.py => lightning_extension.py} (74%) rename python/nano/test/{test_models_lightning_support.py => test_models_lightning_extension.py} (90%) diff --git a/python/nano/src/bigdl/nano/pytorch/vision/models/lightning_support.py b/python/nano/src/bigdl/nano/pytorch/vision/models/lightning_extension.py similarity index 74% rename from python/nano/src/bigdl/nano/pytorch/vision/models/lightning_support.py rename to python/nano/src/bigdl/nano/pytorch/vision/models/lightning_extension.py index 72ded4158c1..228284c4970 100644 --- a/python/nano/src/bigdl/nano/pytorch/vision/models/lightning_support.py +++ b/python/nano/src/bigdl/nano/pytorch/vision/models/lightning_extension.py @@ -13,15 +13,27 @@ # See the License for the specific language governing permissions and # limitations under the License. # +from typing import Callable import pytorch_lightning as pl -def lightning_module(loss_creator, optimizer_creator, configs): - def lightning_from_torch(cls): +def to_lightning(loss_creator: Callable, optimizer_creator: Callable, configs: dict): + """ + A decorator on torch module creator, returns a pytorch-lightning model. + Args: + loss_creator: A function takes configs in and returns loss class. + optimizer_creator: A function takes torch module and configs in, returns optimizer. + configs: A dict with key-values for loss and optimizer creator. + + + Returns: Decorator function on class or function + + """ + def from_torch(creator): class LightningModel(pl.LightningModule): def __init__(self, *args, **kwargs): super().__init__() - self.model = cls(*args, **kwargs) + self.model = creator(*args, **kwargs) self.loss = loss_creator(configs) self.optimizer = optimizer_creator(self.model, configs) @@ -53,4 +65,4 @@ def configure_optimizers(self): return LightningModel - return lightning_from_torch + return from_torch diff --git a/python/nano/test/test_models_lightning_support.py b/python/nano/test/test_models_lightning_extension.py similarity index 90% rename from python/nano/test/test_models_lightning_support.py rename to python/nano/test/test_models_lightning_extension.py index 3717b6421fd..e88ab7fb26a 100644 --- a/python/nano/test/test_models_lightning_support.py +++ b/python/nano/test/test_models_lightning_extension.py @@ -21,7 +21,8 @@ from test._train_torch_lightning import train_torch_lightning from torch import nn -from bigdl.nano.pytorch.vision.models import lightning_support, vision +from bigdl.nano.pytorch.vision.models import vision +from bigdl.nano.pytorch.vision.models.lightning_extension import to_lightning config = { 'lr': 0.01, @@ -40,7 +41,7 @@ def optimizer_creator(model, config): return getattr(torch.optim, config.get("optim", "Adam"))(model.parameters(), lr=config.get("lr", 0.001)) -@lightning_support.lightning_module(loss_creator, optimizer_creator, config) +@to_lightning(loss_creator, optimizer_creator, config) def resnet18(num_classes, pretrained=True, include_top=False, freeze=True): backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) output_size = backbone.get_output_size() From c6f512763e1ac833a0b8f83f37be747b9a35b6d1 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 28 Oct 2021 15:32:55 +0800 Subject: [PATCH 04/30] moved lightning extension to nano/pytorch --- .../nano/pytorch/{vision/models => }/lightning_extension.py | 0 python/nano/test/test_models_lightning_extension.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename python/nano/src/bigdl/nano/pytorch/{vision/models => }/lightning_extension.py (100%) diff --git a/python/nano/src/bigdl/nano/pytorch/vision/models/lightning_extension.py b/python/nano/src/bigdl/nano/pytorch/lightning_extension.py similarity index 100% rename from python/nano/src/bigdl/nano/pytorch/vision/models/lightning_extension.py rename to python/nano/src/bigdl/nano/pytorch/lightning_extension.py diff --git a/python/nano/test/test_models_lightning_extension.py b/python/nano/test/test_models_lightning_extension.py index e88ab7fb26a..1f6ac42f424 100644 --- a/python/nano/test/test_models_lightning_extension.py +++ b/python/nano/test/test_models_lightning_extension.py @@ -22,7 +22,7 @@ from torch import nn from bigdl.nano.pytorch.vision.models import vision -from bigdl.nano.pytorch.vision.models.lightning_extension import to_lightning +from bigdl.nano.pytorch.lightning_extension import to_lightning config = { 'lr': 0.01, From 6580715c5941f9b34f7843a5d11ad33bc6dbd706 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 28 Oct 2021 16:18:17 +0800 Subject: [PATCH 05/30] remove loss, optim creator and directly pass loss and optimizer to initiate --- .../bigdl/nano/pytorch/lightning_extension.py | 18 ++++++++++-------- .../test/test_models_lightning_extension.py | 18 ++++-------------- 2 files changed, 14 insertions(+), 22 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning_extension.py b/python/nano/src/bigdl/nano/pytorch/lightning_extension.py index 228284c4970..4310badf913 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning_extension.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning_extension.py @@ -13,29 +13,31 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Callable import pytorch_lightning as pl +import torch +from torch.nn.modules.loss import _Loss -def to_lightning(loss_creator: Callable, optimizer_creator: Callable, configs: dict): +def to_lightning(loss: _Loss, optimizer: torch.optim, **opt_args): """ A decorator on torch module creator, returns a pytorch-lightning model. - Args: - loss_creator: A function takes configs in and returns loss class. - optimizer_creator: A function takes torch module and configs in, returns optimizer. - configs: A dict with key-values for loss and optimizer creator. + Args: + loss: torch loss function. + optimizer: torch optimizer. + **opt_args: arguments for optimizer. Returns: Decorator function on class or function """ + def from_torch(creator): class LightningModel(pl.LightningModule): def __init__(self, *args, **kwargs): super().__init__() self.model = creator(*args, **kwargs) - self.loss = loss_creator(configs) - self.optimizer = optimizer_creator(self.model, configs) + self.loss = loss + self.optimizer = optimizer(self.model.parameters(), **opt_args) def forward(self, batch): # Handle different numbers of input for various models diff --git a/python/nano/test/test_models_lightning_extension.py b/python/nano/test/test_models_lightning_extension.py index 1f6ac42f424..fa098be99d2 100644 --- a/python/nano/test/test_models_lightning_extension.py +++ b/python/nano/test/test_models_lightning_extension.py @@ -21,27 +21,17 @@ from test._train_torch_lightning import train_torch_lightning from torch import nn -from bigdl.nano.pytorch.vision.models import vision from bigdl.nano.pytorch.lightning_extension import to_lightning +from bigdl.nano.pytorch.vision.models import vision -config = { - 'lr': 0.01, - 'optim': 'Adam', -} batch_size = 256 num_workers = 0 data_dir = os.path.join(os.path.dirname(__file__), "data") - -def loss_creator(config): - return nn.CrossEntropyLoss() - - -def optimizer_creator(model, config): - return getattr(torch.optim, config.get("optim", "Adam"))(model.parameters(), lr=config.get("lr", 0.001)) +loss = nn.CrossEntropyLoss() -@to_lightning(loss_creator, optimizer_creator, config) +@to_lightning(loss, torch.optim.Adam, lr=0.01) def resnet18(num_classes, pretrained=True, include_top=False, freeze=True): backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) output_size = backbone.get_output_size() @@ -49,7 +39,7 @@ def resnet18(num_classes, pretrained=True, include_top=False, freeze=True): return torch.nn.Sequential(backbone, head) -class TestModelsLightningSupport(TestCase): +class TestModelsLightningExtension(TestCase): num_classes = 10 def test_resnet18_ipex(self): From 06f4b88092a44405476d317fdff4655fd5d61c8f Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Fri, 29 Oct 2021 16:55:42 +0800 Subject: [PATCH 06/30] added another implementation for pytorch to lightning --- python/nano/test/pytorch_to_lightning_v1.py | 191 ++++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 python/nano/test/pytorch_to_lightning_v1.py diff --git a/python/nano/test/pytorch_to_lightning_v1.py b/python/nano/test/pytorch_to_lightning_v1.py new file mode 100644 index 00000000000..ad580dda098 --- /dev/null +++ b/python/nano/test/pytorch_to_lightning_v1.py @@ -0,0 +1,191 @@ +# +# Copyright 2016 The BigDL Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import os + +import torch +from pytorch_lightning import LightningModule +from torch.nn.modules.loss import _Loss + +from _train_torch_lightning import create_data_loader, data_transform +from bigdl.nano.pytorch.onnxruntime_support import onnxruntime +from bigdl.nano.pytorch.vision.models import vision + + +def to_lightning(loss: _Loss, optimizer: torch.optim, **opt_args): + """ + A decorator on torch module creator, returns a pytorch-lightning model. + + Args: + loss: torch loss function. + optimizer: torch optimizer. + **opt_args: arguments for optimizer. + + Returns: Decorator function on class or function + + """ + + def from_torch(creator): + + class LightningModel(LightningModule): + def __init__(self, *args, **kwargs): + super().__init__() + if not isinstance(creator, nn.Module): + torch_model = creator(*args, **kwargs) + self.copy(torch_model) + + @staticmethod + def from_instance(instance): + pl_model = LightningModel() + pl_model.copy(instance) + return pl_model + + def copy(self, torch_model): + for k in torch_model.__dict__.keys(): + setattr(self, k, getattr(torch_model, k)) + self.forward = torch_model.forward + + @property + def loss(self): + return loss + + def _forward(self, batch): + # Handle different numbers of input for various models + nargs = self.forward.__code__.co_argcount + return self.forward(*(batch[:nargs - 1])) + + def training_step(self, batch, batch_idx): + x, y = batch + y_hat = self._forward(batch) + loss = self.loss(y_hat, y) + return loss + + def validation_step(self, batch, batch_idx): + x, y = batch + y_hat = self._forward(batch) + loss = self.loss(y_hat, y) + return loss + + def test_step(self, batch, batch_idx): + x, y = batch + y_hat = self._forward(batch) + loss = self.loss(y_hat, y) + return loss + + def configure_optimizers(self): + return optimizer(self.parameters(), **opt_args) + + if not isinstance(creator, nn.Module): + return LightningModel + else: + return LightningModel.from_instance(creator) + + return from_torch + + +from torch import nn + +loss = nn.CrossEntropyLoss() +x = torch.ones(1, 3, 256, 256) +y = torch.zeros(1, 254, 254, dtype=torch.long) + + +# Case 1 decorate a module class +@to_lightning(loss, torch.optim.Adam, lr=0.01) +class Net(nn.Module): + def __init__(self, ic, oc): + super().__init__() + self.conv1 = nn.Conv2d(ic, oc, 3) + + # for test + def _forward(self, x): + return self.conv1(x) + + def forward(self, x): + return self._forward(x) + + +model = Net(3, 1) +out = model.training_step([x, y], 0) +print(out.shape) + +batch_size = 256 +num_workers = 0 +data_dir = os.path.join(os.path.dirname(__file__), "data") +from bigdl.nano.pytorch.trainer import Trainer + +trainer = Trainer(max_epochs=1) +data_loader = create_data_loader( + data_dir, batch_size, num_workers, data_transform) + + +# Case 2 convert from instance +def resnet18(num_classes, pretrained=True, include_top=False, freeze=True): + backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) + output_size = backbone.get_output_size() + head = nn.Linear(output_size, num_classes) + return torch.nn.Sequential(backbone, head) + + +model: nn.Module = resnet18(10, pretrained=True, include_top=False, freeze=True) +# Save to test if load_state_dict work after conversion +torch.save(model.state_dict(), "./tmp.pth") + +pl_model1 = (to_lightning(loss, torch.optim.Adam, lr=0.01)(model)) +trainer.fit(pl_model1, data_loader) +trainer.test(pl_model1, data_loader) + + +# Case 3 decorate a function +@to_lightning(loss, torch.optim.Adam, lr=0.01) +def decorated_resnet18(num_classes, pretrained=True, include_top=False, freeze=True): + return resnet18(num_classes, pretrained=pretrained, include_top=include_top, freeze=freeze) + + +pl_model2 = decorated_resnet18(10, pretrained=True, include_top=False, freeze=True) +# trainer.fit(pl_model2, data_loader) +# trainer.test(pl_model2, data_loader) + +# Test if pl_model can load saved keys by torch +pl_model2.load_state_dict(torch.load('tmp.pth'), strict=True) + + +def composed(*decs): + def deco(f): + for dec in reversed(decs): + f = dec(f) + return f + return deco + +def preprocess(loss=None, optimizer=None, config=None, onnx=True): + return composed( + onnxruntime(onnx), + to_lightning(loss, optimizer, **config) + ) + + +# Case 4 overall decorated +@preprocess(loss, torch.optim.Adam, {"lr": 0.01}, onnx=True) +def decorated_resnet18(num_classes, pretrained=True, include_top=False, freeze=True): + return resnet18(num_classes, pretrained=pretrained, include_top=include_top, freeze=freeze) + +pl_model3 = decorated_resnet18(10, pretrained=True, include_top=False, freeze=True) +# trainer.fit(pl_model3, data_loader) +# trainer.test(pl_model3, data_loader) + +# Need to modify onnxruntime decorator +pl_model4 = preprocess(loss, torch.optim.Adam, {"lr": 0.01}, onnx=True)(model) +# trainer.fit(pl_model4, data_loader) +# trainer.test(pl_model4, data_loader) \ No newline at end of file From ff2f028c566f814b24fac7f85b18feacf0fb4124 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 3 Nov 2021 13:16:45 +0800 Subject: [PATCH 07/30] use LightningModuleFromTorch to create lightning module from pytorch --- .../nano/src/bigdl/nano/pytorch/lightning.py | 74 +++++++++++++++++++ .../bigdl/nano/pytorch/lightning_extension.py | 70 ------------------ ...ghtning_extension.py => test_lightning.py} | 27 ++++--- 3 files changed, 91 insertions(+), 80 deletions(-) create mode 100644 python/nano/src/bigdl/nano/pytorch/lightning.py delete mode 100644 python/nano/src/bigdl/nano/pytorch/lightning_extension.py rename python/nano/test/{test_models_lightning_extension.py => test_lightning.py} (53%) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py new file mode 100644 index 00000000000..9a9b9761a79 --- /dev/null +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -0,0 +1,74 @@ +# +# Copyright 2016 The BigDL Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from pytorch_lightning import LightningModule +from torch import nn +import torch +from torch.nn.modules.loss import _Loss + + +class LightningModuleFromTorch(LightningModule): + def __init__(self, model: nn.Module, loss: _Loss, optimizer: torch.optim): + """ + Integrate pytorch modules, loss, optimizer to pytorch-lightning model. + + Args: + model: pytorch model to be converted. + loss: torch loss function. + optimizer: torch optimizer. + + Returns: LightningModule Object + """ + super().__init__() + self.copy(model) + self._loss = loss + self._optimizer = optimizer + + def copy(self, torch_model): + for name, child in torch_model._modules.items(): + setattr(self, name, child) + self.forward = torch_model.forward + + @property + def loss(self): + return self._loss + + def _forward(self, batch): + # Handle different numbers of input for various models + nargs = self.forward.__code__.co_argcount + return self.forward(*(batch[:nargs - 1])) + + def training_step(self, batch, batch_idx): + x, y = batch + y_hat = self._forward(batch) + loss = self.loss(y_hat, y) + return loss + + def validation_step(self, batch, batch_idx): + x, y = batch + y_hat = self._forward(batch) + loss = self.loss(y_hat, y) + return loss + + def test_step(self, batch, batch_idx): + x, y = batch + y_hat = self._forward(batch) + loss = self.loss(y_hat, y) + return loss + + def configure_optimizers(self): + return self._optimizer + + diff --git a/python/nano/src/bigdl/nano/pytorch/lightning_extension.py b/python/nano/src/bigdl/nano/pytorch/lightning_extension.py deleted file mode 100644 index 4310badf913..00000000000 --- a/python/nano/src/bigdl/nano/pytorch/lightning_extension.py +++ /dev/null @@ -1,70 +0,0 @@ -# -# Copyright 2016 The BigDL Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import pytorch_lightning as pl -import torch -from torch.nn.modules.loss import _Loss - - -def to_lightning(loss: _Loss, optimizer: torch.optim, **opt_args): - """ - A decorator on torch module creator, returns a pytorch-lightning model. - - Args: - loss: torch loss function. - optimizer: torch optimizer. - **opt_args: arguments for optimizer. - - Returns: Decorator function on class or function - - """ - - def from_torch(creator): - class LightningModel(pl.LightningModule): - def __init__(self, *args, **kwargs): - super().__init__() - self.model = creator(*args, **kwargs) - self.loss = loss - self.optimizer = optimizer(self.model.parameters(), **opt_args) - - def forward(self, batch): - # Handle different numbers of input for various models - nargs = self.model.forward.__code__.co_argcount - return self.model(*batch[:nargs - 1]) - - def training_step(self, batch, batch_idx): - x, y = batch - y_hat = self(batch) - loss = self.loss(y_hat, y) - return loss - - def validation_step(self, batch, batch_idx): - x, y = batch - y_hat = self(batch) - loss = self.loss(y_hat, y) - return loss - - def test_step(self, batch, batch_idx): - x, y = batch - y_hat = self(batch) - loss = self.loss(y_hat, y) - return loss - - def configure_optimizers(self): - return self.optimizer - - return LightningModel - - return from_torch diff --git a/python/nano/test/test_models_lightning_extension.py b/python/nano/test/test_lightning.py similarity index 53% rename from python/nano/test/test_models_lightning_extension.py rename to python/nano/test/test_lightning.py index fa098be99d2..bbebe333711 100644 --- a/python/nano/test/test_models_lightning_extension.py +++ b/python/nano/test/test_lightning.py @@ -21,7 +21,7 @@ from test._train_torch_lightning import train_torch_lightning from torch import nn -from bigdl.nano.pytorch.lightning_extension import to_lightning +from bigdl.nano.pytorch.lightning import LightningModuleFromTorch from bigdl.nano.pytorch.vision.models import vision batch_size = 256 @@ -31,18 +31,25 @@ loss = nn.CrossEntropyLoss() -@to_lightning(loss, torch.optim.Adam, lr=0.01) -def resnet18(num_classes, pretrained=True, include_top=False, freeze=True): - backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) - output_size = backbone.get_output_size() - head = nn.Linear(output_size, num_classes) - return torch.nn.Sequential(backbone, head) +class ResNet18(nn.Module): + def __init__(self, num_classes, pretrained=True, include_top=False, freeze=True): + super().__init__() + backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) + output_size = backbone.get_output_size() + head = nn.Linear(output_size, num_classes) + self.model = torch.nn.Sequential(backbone, head) + def forward(self, x): + return self.model(x) -class TestModelsLightningExtension(TestCase): + +class TestLightningModuleFromTorch(TestCase): num_classes = 10 def test_resnet18_ipex(self): - pl_resnet18 = resnet18(10, pretrained=True, include_top=False, freeze=True) - train_torch_lightning(pl_resnet18, batch_size, num_workers, data_dir, + model = ResNet18(10, pretrained=True, include_top=False, freeze=True) + loss = nn.CrossEntropyLoss() + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + pl_model = LightningModuleFromTorch(model, loss, optimizer) + train_torch_lightning(pl_model, batch_size, num_workers, data_dir, use_orca_lite_trainer=True) From aae8fe9f9eb26f410eb7d96404b9bf23bc8c8c3e Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 3 Nov 2021 13:17:32 +0800 Subject: [PATCH 08/30] remove temporary change --- python/nano/test/pytorch_to_lightning_v1.py | 191 -------------------- 1 file changed, 191 deletions(-) delete mode 100644 python/nano/test/pytorch_to_lightning_v1.py diff --git a/python/nano/test/pytorch_to_lightning_v1.py b/python/nano/test/pytorch_to_lightning_v1.py deleted file mode 100644 index ad580dda098..00000000000 --- a/python/nano/test/pytorch_to_lightning_v1.py +++ /dev/null @@ -1,191 +0,0 @@ -# -# Copyright 2016 The BigDL Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import os - -import torch -from pytorch_lightning import LightningModule -from torch.nn.modules.loss import _Loss - -from _train_torch_lightning import create_data_loader, data_transform -from bigdl.nano.pytorch.onnxruntime_support import onnxruntime -from bigdl.nano.pytorch.vision.models import vision - - -def to_lightning(loss: _Loss, optimizer: torch.optim, **opt_args): - """ - A decorator on torch module creator, returns a pytorch-lightning model. - - Args: - loss: torch loss function. - optimizer: torch optimizer. - **opt_args: arguments for optimizer. - - Returns: Decorator function on class or function - - """ - - def from_torch(creator): - - class LightningModel(LightningModule): - def __init__(self, *args, **kwargs): - super().__init__() - if not isinstance(creator, nn.Module): - torch_model = creator(*args, **kwargs) - self.copy(torch_model) - - @staticmethod - def from_instance(instance): - pl_model = LightningModel() - pl_model.copy(instance) - return pl_model - - def copy(self, torch_model): - for k in torch_model.__dict__.keys(): - setattr(self, k, getattr(torch_model, k)) - self.forward = torch_model.forward - - @property - def loss(self): - return loss - - def _forward(self, batch): - # Handle different numbers of input for various models - nargs = self.forward.__code__.co_argcount - return self.forward(*(batch[:nargs - 1])) - - def training_step(self, batch, batch_idx): - x, y = batch - y_hat = self._forward(batch) - loss = self.loss(y_hat, y) - return loss - - def validation_step(self, batch, batch_idx): - x, y = batch - y_hat = self._forward(batch) - loss = self.loss(y_hat, y) - return loss - - def test_step(self, batch, batch_idx): - x, y = batch - y_hat = self._forward(batch) - loss = self.loss(y_hat, y) - return loss - - def configure_optimizers(self): - return optimizer(self.parameters(), **opt_args) - - if not isinstance(creator, nn.Module): - return LightningModel - else: - return LightningModel.from_instance(creator) - - return from_torch - - -from torch import nn - -loss = nn.CrossEntropyLoss() -x = torch.ones(1, 3, 256, 256) -y = torch.zeros(1, 254, 254, dtype=torch.long) - - -# Case 1 decorate a module class -@to_lightning(loss, torch.optim.Adam, lr=0.01) -class Net(nn.Module): - def __init__(self, ic, oc): - super().__init__() - self.conv1 = nn.Conv2d(ic, oc, 3) - - # for test - def _forward(self, x): - return self.conv1(x) - - def forward(self, x): - return self._forward(x) - - -model = Net(3, 1) -out = model.training_step([x, y], 0) -print(out.shape) - -batch_size = 256 -num_workers = 0 -data_dir = os.path.join(os.path.dirname(__file__), "data") -from bigdl.nano.pytorch.trainer import Trainer - -trainer = Trainer(max_epochs=1) -data_loader = create_data_loader( - data_dir, batch_size, num_workers, data_transform) - - -# Case 2 convert from instance -def resnet18(num_classes, pretrained=True, include_top=False, freeze=True): - backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) - output_size = backbone.get_output_size() - head = nn.Linear(output_size, num_classes) - return torch.nn.Sequential(backbone, head) - - -model: nn.Module = resnet18(10, pretrained=True, include_top=False, freeze=True) -# Save to test if load_state_dict work after conversion -torch.save(model.state_dict(), "./tmp.pth") - -pl_model1 = (to_lightning(loss, torch.optim.Adam, lr=0.01)(model)) -trainer.fit(pl_model1, data_loader) -trainer.test(pl_model1, data_loader) - - -# Case 3 decorate a function -@to_lightning(loss, torch.optim.Adam, lr=0.01) -def decorated_resnet18(num_classes, pretrained=True, include_top=False, freeze=True): - return resnet18(num_classes, pretrained=pretrained, include_top=include_top, freeze=freeze) - - -pl_model2 = decorated_resnet18(10, pretrained=True, include_top=False, freeze=True) -# trainer.fit(pl_model2, data_loader) -# trainer.test(pl_model2, data_loader) - -# Test if pl_model can load saved keys by torch -pl_model2.load_state_dict(torch.load('tmp.pth'), strict=True) - - -def composed(*decs): - def deco(f): - for dec in reversed(decs): - f = dec(f) - return f - return deco - -def preprocess(loss=None, optimizer=None, config=None, onnx=True): - return composed( - onnxruntime(onnx), - to_lightning(loss, optimizer, **config) - ) - - -# Case 4 overall decorated -@preprocess(loss, torch.optim.Adam, {"lr": 0.01}, onnx=True) -def decorated_resnet18(num_classes, pretrained=True, include_top=False, freeze=True): - return resnet18(num_classes, pretrained=pretrained, include_top=include_top, freeze=freeze) - -pl_model3 = decorated_resnet18(10, pretrained=True, include_top=False, freeze=True) -# trainer.fit(pl_model3, data_loader) -# trainer.test(pl_model3, data_loader) - -# Need to modify onnxruntime decorator -pl_model4 = preprocess(loss, torch.optim.Adam, {"lr": 0.01}, onnx=True)(model) -# trainer.fit(pl_model4, data_loader) -# trainer.test(pl_model4, data_loader) \ No newline at end of file From 6af75353b5de60a8502fbfe8fb928bf4806d0717 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 3 Nov 2021 13:24:45 +0800 Subject: [PATCH 09/30] remove redundant part --- python/nano/src/bigdl/nano/pytorch/lightning.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index 9a9b9761a79..c6e61bf4a67 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -33,18 +33,14 @@ def __init__(self, model: nn.Module, loss: _Loss, optimizer: torch.optim): """ super().__init__() self.copy(model) - self._loss = loss - self._optimizer = optimizer + self.loss = loss + self.optimizer = optimizer def copy(self, torch_model): for name, child in torch_model._modules.items(): setattr(self, name, child) self.forward = torch_model.forward - @property - def loss(self): - return self._loss - def _forward(self, batch): # Handle different numbers of input for various models nargs = self.forward.__code__.co_argcount @@ -69,6 +65,6 @@ def test_step(self, batch, batch_idx): return loss def configure_optimizers(self): - return self._optimizer + return self.optimizer From c02c13ed3597f2b98dd2ede7109ce1bf3ce66539 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 3 Nov 2021 13:28:45 +0800 Subject: [PATCH 10/30] added trainer.compile to convert pytorch to pytorch-lightning --- .../src/bigdl/nano/pytorch/trainer/Trainer.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index 632c82064f0..8674bc3b967 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -18,6 +18,10 @@ from logging import warning import torch import pytorch_lightning as pl +from torch import nn +from torch.nn.modules.loss import _Loss + +from bigdl.nano.pytorch.lightning import LightningModuleFromTorch from bigdl.nano.pytorch.plugins.ddp_spawn import DDPSpawnPlugin from bigdl.nano.common import check_avx512 from pytorch_lightning.plugins.environments import LightningEnvironment @@ -104,3 +108,23 @@ def __init__(self, num_processes: int = 1, super().__init__(accelerator=accelerator, plugins=[plugin], *args, **kwargs) + + @staticmethod + def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None): + """ + Compile a pytorch model into a pytorch-lightning model and return it. + Args: + model: A pytorch model instance. + loss: Loss for LightningModule. + optimizer: Optimizer for LightningModule. + + Returns: A LightningModule converted from model. + + """ + if isinstance(model, pl.LightningModule): + if loss: + model.loss = loss + if optimizer: + model.optimizer = optimizer + else: + return LightningModuleFromTorch(model, loss, optimizer) From 4b930f23b541570e3c4cabd94d4113f83a015fef Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 3 Nov 2021 13:42:36 +0800 Subject: [PATCH 11/30] added unit test for trainer.compile --- python/nano/test/test_trainer_ipex.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/python/nano/test/test_trainer_ipex.py b/python/nano/test/test_trainer_ipex.py index a4104bf7693..f865a744996 100644 --- a/python/nano/test/test_trainer_ipex.py +++ b/python/nano/test/test_trainer_ipex.py @@ -18,10 +18,15 @@ import pytest import os from unittest import TestCase + +import torch +from torch import nn + +from _train_torch_lightning import train_torch_lightning +from bigdl.nano.pytorch.trainer import Trainer from bigdl.nano.pytorch.vision.models import vision from test._train_torch_lightning import train_with_linear_top_layer - batch_size = 256 num_workers = 0 data_dir = os.path.join(os.path.dirname(__file__), "data") @@ -78,6 +83,25 @@ def test_shufflenet_ipex(self): shufflenet, batch_size, num_workers, data_dir, use_orca_lite_trainer=True) + def test_trainer_compile(self): + class ResNet18(nn.Module): + def __init__(self, num_classes, pretrained=True, include_top=False, freeze=True): + super().__init__() + backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) + output_size = backbone.get_output_size() + head = nn.Linear(output_size, num_classes) + self.model = nn.Sequential(backbone, head) + + def forward(self, x): + return self.model(x) + + model = ResNet18(10, pretrained=True, include_top=False, freeze=True) + loss = nn.CrossEntropyLoss() + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + pl_model = Trainer.compile(model, loss, optimizer) + train_torch_lightning(pl_model, batch_size, num_workers, data_dir, + use_orca_lite_trainer=True) + if __name__ == '__main__': pytest.main([__file__]) From 6aa2dc187b68bcfe82f8c04bee0951e83f5e450f Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 3 Nov 2021 14:04:33 +0800 Subject: [PATCH 12/30] fixed return when input is pl model --- python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index 8674bc3b967..3625022499b 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -126,5 +126,6 @@ def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None) model.loss = loss if optimizer: model.optimizer = optimizer + return model else: return LightningModuleFromTorch(model, loss, optimizer) From d5bb5c1d1300bb85a5366988ff030911a417f7ac Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 3 Nov 2021 14:08:17 +0800 Subject: [PATCH 13/30] added type hint for LightningModuleFromTorch.copy --- python/nano/src/bigdl/nano/pytorch/lightning.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index c6e61bf4a67..c0790252a6c 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -36,7 +36,7 @@ def __init__(self, model: nn.Module, loss: _Loss, optimizer: torch.optim): self.loss = loss self.optimizer = optimizer - def copy(self, torch_model): + def copy(self, torch_model: nn.Module): for name, child in torch_model._modules.items(): setattr(self, name, child) self.forward = torch_model.forward From a893aa313b9db1fc14bd93b872f168ca81ffb890 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 09:09:42 +0800 Subject: [PATCH 14/30] Renamed copy as _copy --- python/nano/src/bigdl/nano/pytorch/lightning.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index c0790252a6c..78458849bcd 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -32,11 +32,11 @@ def __init__(self, model: nn.Module, loss: _Loss, optimizer: torch.optim): Returns: LightningModule Object """ super().__init__() - self.copy(model) + self._copy(model) self.loss = loss self.optimizer = optimizer - def copy(self, torch_model: nn.Module): + def _copy(self, torch_model: nn.Module): for name, child in torch_model._modules.items(): setattr(self, name, child) self.forward = torch_model.forward From c6fb6936dccd613c1dadacf416e8cd7ef9059351 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 09:28:19 +0800 Subject: [PATCH 15/30] Modified comment of compile --- .../nano/src/bigdl/nano/pytorch/trainer/Trainer.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index 3625022499b..3420b507ec7 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -112,13 +112,15 @@ def __init__(self, num_processes: int = 1, @staticmethod def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None): """ - Compile a pytorch model into a pytorch-lightning model and return it. + Construct a pytorch-lightning model. If model is already a pytorch-lightning model, return model. + If model is pytorch model, construct a new pytorch-lightning module with model, loss and optimizer. + Args: - model: A pytorch model instance. - loss: Loss for LightningModule. - optimizer: Optimizer for LightningModule. + model: A model instance. + loss: Loss to construct pytorch-lightning model. Should be None if model is instance of pl.LightningModule. + optimizer: Optimizer to construct pytorch-lightning model Should be None if model is instance of pl.LightningModule. - Returns: A LightningModule converted from model. + Returns: A LightningModule object. """ if isinstance(model, pl.LightningModule): From 596e4e6b45df686353a58d174dd95beed325622f Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 10:08:30 +0800 Subject: [PATCH 16/30] added input checking --- .../src/bigdl/nano/pytorch/trainer/Trainer.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index 3420b507ec7..697e36c72d0 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -16,16 +16,17 @@ from logging import warning -import torch +from typing import Any, List, Optional + import pytorch_lightning as pl +import torch +from pytorch_lightning.plugins.environments import LightningEnvironment from torch import nn from torch.nn.modules.loss import _Loss +from bigdl.nano.common import check_avx512 from bigdl.nano.pytorch.lightning import LightningModuleFromTorch from bigdl.nano.pytorch.plugins.ddp_spawn import DDPSpawnPlugin -from bigdl.nano.common import check_avx512 -from pytorch_lightning.plugins.environments import LightningEnvironment -from typing import Any, List, Optional distributed_backends = ["spawn", "ray"] @@ -123,11 +124,10 @@ def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None) Returns: A LightningModule object. """ + assert isinstance(model, nn.Module), "Model must be instance of nn.Module but got {}".format(model.__class__) if isinstance(model, pl.LightningModule): - if loss: - model.loss = loss - if optimizer: - model.optimizer = optimizer + assert not (loss or optimizer), \ + "Loss and optimizer should be None if model is already a pytorch-lightning model." return model else: return LightningModuleFromTorch(model, loss, optimizer) From db4466b5b5b2e8fce1ea6f9805dcb814bbf054c8 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 10:25:36 +0800 Subject: [PATCH 17/30] refactored docstring --- python/nano/src/bigdl/nano/pytorch/lightning.py | 9 +++------ .../nano/src/bigdl/nano/pytorch/trainer/Trainer.py | 13 ++++++------- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index 78458849bcd..7c2f5911077 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -24,12 +24,9 @@ def __init__(self, model: nn.Module, loss: _Loss, optimizer: torch.optim): """ Integrate pytorch modules, loss, optimizer to pytorch-lightning model. - Args: - model: pytorch model to be converted. - loss: torch loss function. - optimizer: torch optimizer. - - Returns: LightningModule Object + :param model: pytorch model to be converted. + :param loss: torch loss function. + :param optimizer: torch optimizer. """ super().__init__() self._copy(model) diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index 697e36c72d0..e87aadaaa54 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -116,13 +116,12 @@ def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None) Construct a pytorch-lightning model. If model is already a pytorch-lightning model, return model. If model is pytorch model, construct a new pytorch-lightning module with model, loss and optimizer. - Args: - model: A model instance. - loss: Loss to construct pytorch-lightning model. Should be None if model is instance of pl.LightningModule. - optimizer: Optimizer to construct pytorch-lightning model Should be None if model is instance of pl.LightningModule. - - Returns: A LightningModule object. - + :param model: A model instance. + :param loss: Loss to construct pytorch-lightning model. Should be None if model is instance of + pl.LightningModule. + :param optimizer: Optimizer to construct pytorch-lightning model Should be None if model is instance + of pl.LightningModule. + :return: A LightningModule object. """ assert isinstance(model, nn.Module), "Model must be instance of nn.Module but got {}".format(model.__class__) if isinstance(model, pl.LightningModule): From 90eabc7700278cf7d65c29ba6fc7c46ffe65fc9d Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 10:59:18 +0800 Subject: [PATCH 18/30] Reformat docstring --- python/nano/src/bigdl/nano/pytorch/lightning.py | 6 +++--- .../nano/src/bigdl/nano/pytorch/trainer/Trainer.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index 7c2f5911077..419123cadc4 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -24,9 +24,9 @@ def __init__(self, model: nn.Module, loss: _Loss, optimizer: torch.optim): """ Integrate pytorch modules, loss, optimizer to pytorch-lightning model. - :param model: pytorch model to be converted. - :param loss: torch loss function. - :param optimizer: torch optimizer. + :param model: Pytorch model to be converted. + :param loss: A torch loss function. + :param optimizer: A torch optimizer. """ super().__init__() self._copy(model) diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index e87aadaaa54..58cfbe54422 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -116,12 +116,12 @@ def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None) Construct a pytorch-lightning model. If model is already a pytorch-lightning model, return model. If model is pytorch model, construct a new pytorch-lightning module with model, loss and optimizer. - :param model: A model instance. - :param loss: Loss to construct pytorch-lightning model. Should be None if model is instance of - pl.LightningModule. - :param optimizer: Optimizer to construct pytorch-lightning model Should be None if model is instance - of pl.LightningModule. - :return: A LightningModule object. + :param model: A model instance. + :param loss: Loss to construct pytorch-lightning model. Should be None if model is instance of + pl.LightningModule. + :param optimizer: Optimizer to construct pytorch-lightning model Should be None if model is instance of + pl.LightningModule. + :return: A LightningModule object. """ assert isinstance(model, nn.Module), "Model must be instance of nn.Module but got {}".format(model.__class__) if isinstance(model, pl.LightningModule): From 153c1a2a946adc793f4a0b8d9dec3caa18c5125b Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 11:06:06 +0800 Subject: [PATCH 19/30] Tiny changes --- python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index 58cfbe54422..a159046ad78 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -119,14 +119,14 @@ def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None) :param model: A model instance. :param loss: Loss to construct pytorch-lightning model. Should be None if model is instance of pl.LightningModule. - :param optimizer: Optimizer to construct pytorch-lightning model Should be None if model is instance of + :param optimizer: Optimizer to construct pytorch-lightning model Should be None. if model is instance of pl.LightningModule. :return: A LightningModule object. """ assert isinstance(model, nn.Module), "Model must be instance of nn.Module but got {}".format(model.__class__) if isinstance(model, pl.LightningModule): assert not (loss or optimizer), \ - "Loss and optimizer should be None if model is already a pytorch-lightning model." + "Loss and optimizer should be None if model is a pytorch-lightning model." return model else: return LightningModuleFromTorch(model, loss, optimizer) From c642dc9397c93eb9c8df96de5fdd41cb50cd7ece Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 17:48:04 +0800 Subject: [PATCH 20/30] reformat --- python/nano/src/bigdl/nano/pytorch/lightning.py | 2 -- .../src/bigdl/nano/pytorch/trainer/Trainer.py | 16 +++++++++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index 419123cadc4..8905c4ef578 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -63,5 +63,3 @@ def test_step(self, batch, batch_idx): def configure_optimizers(self): return self.optimizer - - diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index a159046ad78..7b56cb0b1ed 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -113,17 +113,19 @@ def __init__(self, num_processes: int = 1, @staticmethod def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None): """ - Construct a pytorch-lightning model. If model is already a pytorch-lightning model, return model. - If model is pytorch model, construct a new pytorch-lightning module with model, loss and optimizer. + Construct a pytorch-lightning model. If model is already a pytorch-lightning model, + return model. If model is pytorch model, construct a new pytorch-lightning module + with model, loss and optimizer. :param model: A model instance. - :param loss: Loss to construct pytorch-lightning model. Should be None if model is instance of - pl.LightningModule. - :param optimizer: Optimizer to construct pytorch-lightning model Should be None. if model is instance of - pl.LightningModule. + :param loss: Loss to construct pytorch-lightning model. + Should be None if model is instance of pl.LightningModule. + :param optimizer: Optimizer to construct pytorch-lightning model Should be None. + if model is instance of pl.LightningModule. :return: A LightningModule object. """ - assert isinstance(model, nn.Module), "Model must be instance of nn.Module but got {}".format(model.__class__) + assert isinstance(model, nn.Module), \ + "Model must be instance of nn.Module but got {}".format(model.__class__) if isinstance(model, pl.LightningModule): assert not (loss or optimizer), \ "Loss and optimizer should be None if model is a pytorch-lightning model." From 1caa2f165c973a7c4a8436a4dc0c393f6288643e Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 18:01:09 +0800 Subject: [PATCH 21/30] correct the import --- python/nano/test/test_trainer_ipex.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/nano/test/test_trainer_ipex.py b/python/nano/test/test_trainer_ipex.py index f865a744996..5b44c7e6c1a 100644 --- a/python/nano/test/test_trainer_ipex.py +++ b/python/nano/test/test_trainer_ipex.py @@ -22,7 +22,7 @@ import torch from torch import nn -from _train_torch_lightning import train_torch_lightning +from test._train_torch_lightning import train_torch_lightning from bigdl.nano.pytorch.trainer import Trainer from bigdl.nano.pytorch.vision.models import vision from test._train_torch_lightning import train_with_linear_top_layer From f38dff282c7a75f3881a90fbc4fafe8ed3311f31 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 4 Nov 2021 18:09:44 +0800 Subject: [PATCH 22/30] type check and --- python/nano/src/bigdl/nano/pytorch/lightning.py | 5 +++-- python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index 8905c4ef578..c0e7fb0ecca 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -17,10 +17,11 @@ from torch import nn import torch from torch.nn.modules.loss import _Loss +from torch.optim import Optimizer class LightningModuleFromTorch(LightningModule): - def __init__(self, model: nn.Module, loss: _Loss, optimizer: torch.optim): + def __init__(self, model: nn.Module, loss: _Loss, optimizer: Optimizer): """ Integrate pytorch modules, loss, optimizer to pytorch-lightning model. @@ -36,7 +37,7 @@ def __init__(self, model: nn.Module, loss: _Loss, optimizer: torch.optim): def _copy(self, torch_model: nn.Module): for name, child in torch_model._modules.items(): setattr(self, name, child) - self.forward = torch_model.forward + setattr(self, "forward", torch_model.forward) def _forward(self, batch): # Handle different numbers of input for various models diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py index 7b56cb0b1ed..ba62be8bd4f 100644 --- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py +++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py @@ -111,7 +111,7 @@ def __init__(self, num_processes: int = 1, plugins=[plugin], *args, **kwargs) @staticmethod - def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None): + def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim.Optimizer = None): """ Construct a pytorch-lightning model. If model is already a pytorch-lightning model, return model. If model is pytorch model, construct a new pytorch-lightning module @@ -131,4 +131,6 @@ def compile(model: nn.Module, loss: _Loss = None, optimizer: torch.optim = None) "Loss and optimizer should be None if model is a pytorch-lightning model." return model else: + assert loss and optimizer, \ + "Loss and optimizer are required to construct a LightningModule instance." return LightningModuleFromTorch(model, loss, optimizer) From 9dbd8f050c98a07550db9906122302bbd9ac3191 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Tue, 9 Nov 2021 09:57:40 +0800 Subject: [PATCH 23/30] assign model as a member variable --- python/nano/src/bigdl/nano/pytorch/lightning.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index c0e7fb0ecca..1557884c4d9 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -30,19 +30,14 @@ def __init__(self, model: nn.Module, loss: _Loss, optimizer: Optimizer): :param optimizer: A torch optimizer. """ super().__init__() - self._copy(model) + self.model = model self.loss = loss self.optimizer = optimizer - def _copy(self, torch_model: nn.Module): - for name, child in torch_model._modules.items(): - setattr(self, name, child) - setattr(self, "forward", torch_model.forward) - def _forward(self, batch): # Handle different numbers of input for various models - nargs = self.forward.__code__.co_argcount - return self.forward(*(batch[:nargs - 1])) + nargs = self.model.forward.__code__.co_argcount + return self.model(*(batch[:nargs-1])) def training_step(self, batch, batch_idx): x, y = batch From 293e54aedc6cc3ab085e1b1ca8391e64ee67e0e7 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Tue, 9 Nov 2021 10:00:11 +0800 Subject: [PATCH 24/30] override load_state_dict --- python/nano/src/bigdl/nano/pytorch/lightning.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index 1557884c4d9..cf2acef7924 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -14,10 +14,11 @@ # limitations under the License. # from pytorch_lightning import LightningModule -from torch import nn +from torch import nn, Tensor import torch from torch.nn.modules.loss import _Loss from torch.optim import Optimizer +from typing import Union, Dict class LightningModuleFromTorch(LightningModule): @@ -59,3 +60,10 @@ def test_step(self, batch, batch_idx): def configure_optimizers(self): return self.optimizer + + def load_state_dict(self, state_dict: Union[Dict[str, Tensor], Dict[str, Tensor]], + strict: bool = True): + try: + self.model.load_state_dict(state_dict) + except RuntimeError: + super().load_state_dict(state_dict) \ No newline at end of file From d3c20d53af868ad66811af40134a53b1de7114ad Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Tue, 9 Nov 2021 10:03:28 +0800 Subject: [PATCH 25/30] fix test_trainer_compile --- python/nano/test/test_trainer_ipex.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/nano/test/test_trainer_ipex.py b/python/nano/test/test_trainer_ipex.py index 5b44c7e6c1a..f19ab4dd0ad 100644 --- a/python/nano/test/test_trainer_ipex.py +++ b/python/nano/test/test_trainer_ipex.py @@ -22,7 +22,7 @@ import torch from torch import nn -from test._train_torch_lightning import train_torch_lightning +from test._train_torch_lightning import create_data_loader, data_transform from bigdl.nano.pytorch.trainer import Trainer from bigdl.nano.pytorch.vision.models import vision from test._train_torch_lightning import train_with_linear_top_layer @@ -98,9 +98,10 @@ def forward(self, x): model = ResNet18(10, pretrained=True, include_top=False, freeze=True) loss = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.01) - pl_model = Trainer.compile(model, loss, optimizer) - train_torch_lightning(pl_model, batch_size, num_workers, data_dir, - use_orca_lite_trainer=True) + trainer = Trainer(max_epochs=1) + pl_model = trainer.compile(model, loss, optimizer) + train_loader = create_data_loader(data_dir, batch_size, num_workers, data_transform) + trainer.fit(pl_model, train_loader) if __name__ == '__main__': From 6fabec3a8ef95ab1496efca7887b95e0d9d7a873 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Tue, 9 Nov 2021 10:15:02 +0800 Subject: [PATCH 26/30] fix test_lightning --- python/nano/test/test_lightning.py | 36 +++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/python/nano/test/test_lightning.py b/python/nano/test/test_lightning.py index bbebe333711..4ac7aa8f688 100644 --- a/python/nano/test/test_lightning.py +++ b/python/nano/test/test_lightning.py @@ -18,21 +18,21 @@ from unittest import TestCase import torch -from test._train_torch_lightning import train_torch_lightning from torch import nn +from _train_torch_lightning import create_data_loader, data_transform from bigdl.nano.pytorch.lightning import LightningModuleFromTorch +from bigdl.nano.pytorch.trainer import Trainer from bigdl.nano.pytorch.vision.models import vision +num_classes = 10 batch_size = 256 num_workers = 0 data_dir = os.path.join(os.path.dirname(__file__), "data") -loss = nn.CrossEntropyLoss() - class ResNet18(nn.Module): - def __init__(self, num_classes, pretrained=True, include_top=False, freeze=True): + def __init__(self, pretrained=True, include_top=False, freeze=True): super().__init__() backbone = vision.resnet18(pretrained=pretrained, include_top=include_top, freeze=freeze) output_size = backbone.get_output_size() @@ -43,13 +43,27 @@ def forward(self, x): return self.model(x) +model = ResNet18(pretrained=True, include_top=False, freeze=True) +loss = nn.CrossEntropyLoss() +optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + + class TestLightningModuleFromTorch(TestCase): - num_classes = 10 - def test_resnet18_ipex(self): - model = ResNet18(10, pretrained=True, include_top=False, freeze=True) - loss = nn.CrossEntropyLoss() - optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + def test_resnet18(self): + pl_model = LightningModuleFromTorch(model, loss, optimizer) + train_loader = create_data_loader(data_dir, batch_size, num_workers, data_transform) + trainer = Trainer(max_epochs=1) + trainer.fit(pl_model, train_loader) + + def test_load_state_dict_from_torch(self): + torch.save(model.state_dict(), "resnet18_test.pth") + pl_model = LightningModuleFromTorch(model, loss, optimizer) + state_dict = torch.load("resnet18_test.pth") + pl_model.load_state_dict(state_dict) + + def test_load_state_dict_from_lightning(self): pl_model = LightningModuleFromTorch(model, loss, optimizer) - train_torch_lightning(pl_model, batch_size, num_workers, data_dir, - use_orca_lite_trainer=True) + torch.save(pl_model.state_dict(), "lightning_resnet18_test.pth") + state_dict = torch.load("lightning_resnet18_test.pth") + pl_model.load_state_dict(state_dict) From d50e403ca6289eff128f96eb562f95b1f1327d47 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Tue, 9 Nov 2021 14:14:01 +0800 Subject: [PATCH 27/30] try lightning module and then self.model --- python/nano/src/bigdl/nano/pytorch/lightning.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index cf2acef7924..38207ce4b3e 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -64,6 +64,6 @@ def configure_optimizers(self): def load_state_dict(self, state_dict: Union[Dict[str, Tensor], Dict[str, Tensor]], strict: bool = True): try: - self.model.load_state_dict(state_dict) + super().load_state_dict(state_dict) except RuntimeError: - super().load_state_dict(state_dict) \ No newline at end of file + self.model.load_state_dict(state_dict) \ No newline at end of file From 66464548d580ec361c1b96ebe905db68545f35f7 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Tue, 9 Nov 2021 14:27:54 +0800 Subject: [PATCH 28/30] rename _forward as forward --- python/nano/src/bigdl/nano/pytorch/lightning.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index 38207ce4b3e..df1d39c10c3 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -13,12 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # +from typing import Union, Dict + from pytorch_lightning import LightningModule from torch import nn, Tensor -import torch from torch.nn.modules.loss import _Loss from torch.optim import Optimizer -from typing import Union, Dict class LightningModuleFromTorch(LightningModule): @@ -35,26 +35,26 @@ def __init__(self, model: nn.Module, loss: _Loss, optimizer: Optimizer): self.loss = loss self.optimizer = optimizer - def _forward(self, batch): + def forward(self, batch): # Handle different numbers of input for various models nargs = self.model.forward.__code__.co_argcount - return self.model(*(batch[:nargs-1])) + return self.model(*(batch[:nargs - 1])) def training_step(self, batch, batch_idx): x, y = batch - y_hat = self._forward(batch) + y_hat = self(batch) loss = self.loss(y_hat, y) return loss def validation_step(self, batch, batch_idx): x, y = batch - y_hat = self._forward(batch) + y_hat = self(batch) loss = self.loss(y_hat, y) return loss def test_step(self, batch, batch_idx): x, y = batch - y_hat = self._forward(batch) + y_hat = self(batch) loss = self.loss(y_hat, y) return loss @@ -66,4 +66,4 @@ def load_state_dict(self, state_dict: Union[Dict[str, Tensor], Dict[str, Tensor] try: super().load_state_dict(state_dict) except RuntimeError: - self.model.load_state_dict(state_dict) \ No newline at end of file + self.model.load_state_dict(state_dict) From 3dc152db2fd8c4e0502e9c030b7247d9d068491c Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Tue, 9 Nov 2021 16:00:57 +0800 Subject: [PATCH 29/30] type check --- python/nano/src/bigdl/nano/pytorch/lightning.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index df1d39c10c3..d9b132ffc52 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # +from collections import OrderedDict from typing import Union, Dict from pytorch_lightning import LightningModule @@ -61,7 +62,7 @@ def test_step(self, batch, batch_idx): def configure_optimizers(self): return self.optimizer - def load_state_dict(self, state_dict: Union[Dict[str, Tensor], Dict[str, Tensor]], + def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]', strict: bool = True): try: super().load_state_dict(state_dict) From 36b8b6ccfe91a8856d8ee85a4c646a6c1305cfc2 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Thu, 11 Nov 2021 14:45:32 +0800 Subject: [PATCH 30/30] optimize imports --- python/nano/src/bigdl/nano/pytorch/lightning.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/nano/src/bigdl/nano/pytorch/lightning.py b/python/nano/src/bigdl/nano/pytorch/lightning.py index d9b132ffc52..ab8242ecb92 100644 --- a/python/nano/src/bigdl/nano/pytorch/lightning.py +++ b/python/nano/src/bigdl/nano/pytorch/lightning.py @@ -14,7 +14,6 @@ # limitations under the License. # from collections import OrderedDict -from typing import Union, Dict from pytorch_lightning import LightningModule from torch import nn, Tensor