Skip to content

[bugfix] Add set_default_tensor_type to torch.DoubleTensor with precision=64 #7108

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Apr 20, 2021
Merged

[bugfix] Add set_default_tensor_type to torch.DoubleTensor with precision=64 #7108

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0ccccee
update
tchaton Apr 19, 2021
0b521e7
Update pytorch_lightning/plugins/precision/double.py
tchaton Apr 20, 2021
1c6f98d
Update pytorch_lightning/plugins/precision/double.py
tchaton Apr 20, 2021
91d5872
Update pytorch_lightning/plugins/precision/double.py
tchaton Apr 20, 2021
be3a5b9
resolve tests
tchaton Apr 20, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pytorch_lightning/accelerators/accelerator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,7 +242,7 @@ def predict_step(self, args: List[Union[Any, int]]) -> STEP_OUTPUT:

args[0] = batch

with self.precision_plugin.predict_context(), self.training_type_plugin.predict_context():
with self.precision_plugin.predict_step_context(), self.training_type_plugin.predict_step_context():
return self.training_type_plugin.predict_step(*args)

def training_step_end(self, output: STEP_OUTPUT) -> STEP_OUTPUT:
Expand Down
2 changes: 1 addition & 1 deletion pytorch_lightning/plugins/base_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,6 @@ def test_step_context(self) -> Generator:
yield

@contextlib.contextmanager
def predict_context(self) -> Generator:
def predict_step_context(self) -> Generator:
"""A contextmanager for the predict step"""
yield
18 changes: 17 additions & 1 deletion pytorch_lightning/plugins/precision/double.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@
# 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 contextlib
from functools import wraps
from typing import Any, List, Tuple
from typing import Any, Generator, List, Tuple

import torch
import torch.nn as nn
Expand Down Expand Up @@ -90,3 +91,18 @@ def connect(
def post_dispatch(self) -> None:
while len(self.patches) > 0:
self.patches.pop().teardown()

@contextlib.contextmanager
def tensor_type_context(self) -> Generator:
"""
A context manager to change the default tensor type.
See: :meth:`torch.set_default_tensor_type`
"""
torch.set_default_tensor_type(torch.DoubleTensor)
yield
torch.set_default_tensor_type(torch.FloatTensor)
Copy link
Collaborator

@Borda Borda Apr 19, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shall you rather save what was it before and then set it back?

Copy link
Contributor Author

@tchaton tchaton Apr 20, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it should be fine. Didn't find any utils to get the default either.


train_step_context = tensor_type_context
val_step_context = tensor_type_context
test_step_context = tensor_type_context
predict_step_context = tensor_type_context
Comment on lines +105 to +108
Copy link
Contributor

@ananthsub ananthsub Apr 21, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tchaton this is breaking the plugin typechecks with mypy:

pytorch_lightning/plugins/precision/double.py:105: error: Incompatible types in assignment (expression has type "Callable[[], _GeneratorContextManager[Any]]", base class "Plugin" defined the type as "Callable[[], Generator[Any, Any, Any]]")  [assignment]
pytorch_lightning/plugins/precision/double.py:106: error: Incompatible types in assignment (expression has type "Callable[[], _GeneratorContextManager[Any]]", base class "Plugin" defined the type as "Callable[[], Generator[Any, Any, Any]]")  [assignment]
pytorch_lightning/plugins/precision/double.py:107: error: Incompatible types in assignment (expression has type "Callable[[], _GeneratorContextManager[Any]]", base class "Plugin" defined the type as "Callable[[], Generator[Any, Any, Any]]")  [assignment]
pytorch_lightning/plugins/precision/double.py:108: error: Incompatible types in assignment (expression has type "Callable[[], _GeneratorContextManager[Any]]", base class "Plugin" defined the type as "Callable[[], Generator[Any, Any, Any]]")  [assignment]

rather than overriding these step contexts, could we share the same implementation?

2 changes: 1 addition & 1 deletion pytorch_lightning/plugins/precision/native_amp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ def test_step_context(self) -> Generator[None, None, None]:
yield

@contextmanager
def predict_context(self) -> Generator[None, None, None]:
def predict_step_context(self) -> Generator[None, None, None]:
"""Enable autocast context"""
with torch.cuda.amp.autocast():
yield
12 changes: 12 additions & 0 deletions tests/plugins/test_double_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,25 +37,37 @@ class DoublePrecisionBoringModel(BoringModel):

def training_step(self, batch, batch_idx):
float_data, int_data = batch
assert torch.tensor([0.]).dtype == torch.float64
assert torch.tensor([0.], dtype=torch.float16).dtype == torch.float16
assert float_data.dtype == torch.float64
output = self(float_data)
loss = self.loss(batch, output)
return {"loss": loss}

def training_epoch_end(self, outputs) -> None:
assert torch.tensor([0.]).dtype == torch.float32
return super().training_epoch_end(outputs)

def validation_step(self, batch, batch_idx):
assert batch.dtype == torch.float64
assert torch.tensor([0.]).dtype == torch.float64
assert torch.tensor([0.], dtype=torch.float16).dtype == torch.float16
output = self(batch)
loss = self.loss(batch, output)
return {"x": loss}

def test_step(self, batch, batch_idx):
assert batch.dtype == torch.float64
assert torch.tensor([0.]).dtype == torch.float64
assert torch.tensor([0.], dtype=torch.float16).dtype == torch.float16
output = self(batch)
loss = self.loss(batch, output)
return {"y": loss}

def predict_step(self, batch, batch_idx, dataloader_idx=None):
assert batch.dtype == torch.float64
assert torch.tensor([0.]).dtype == torch.float64
assert torch.tensor([0.], dtype=torch.float16).dtype == torch.float16
return self(batch)

def on_fit_start(self):
Expand Down