Shortcuts

Source code for pytorch_lightning.accelerators.gpu

# Copyright The PyTorch Lightning team.
#
# 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 logging
import os
import shutil
import subprocess
from typing import Any, Dict, List, Union

import torch

import pytorch_lightning as pl
from pytorch_lightning.accelerators.accelerator import Accelerator
from pytorch_lightning.utilities.exceptions import MisconfigurationException
from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_8

_log = logging.getLogger(__name__)


[docs]class GPUAccelerator(Accelerator): """Accelerator for GPU devices."""
[docs] def setup_environment(self) -> None: """ Raises: MisconfigurationException: If the selected device is not GPU. """ super().setup_environment() if "cuda" not in str(self.training_type_plugin.root_device): raise MisconfigurationException( f"Device should be GPU, got {self.training_type_plugin.root_device} instead" ) torch.cuda.set_device(self.training_type_plugin.root_device)
[docs] def setup(self, trainer: "pl.Trainer") -> None: self.set_nvidia_flags(trainer.local_rank) return super().setup(trainer)
[docs] def on_train_start(self) -> None: super().on_train_start() # clear cache before training torch.cuda.empty_cache()
@staticmethod def set_nvidia_flags(local_rank: int) -> None: # set the correct cuda visible devices (using pci order) os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" all_gpu_ids = ",".join(str(x) for x in range(torch.cuda.device_count())) devices = os.getenv("CUDA_VISIBLE_DEVICES", all_gpu_ids) _log.info(f"LOCAL_RANK: {local_rank} - CUDA_VISIBLE_DEVICES: [{devices}]")
[docs] def get_device_stats(self, device: Union[str, torch.device]) -> Dict[str, Any]: """Gets stats for the given GPU device. Args: device: GPU device for which to get stats Returns: A dictionary mapping the metrics to their values. Raises: FileNotFoundError: If nvidia-smi installation not found """ if _TORCH_GREATER_EQUAL_1_8: return torch.cuda.memory_stats(device) return get_nvidia_gpu_stats(device)
[docs] def teardown(self) -> None: super().teardown() self.training_type_plugin._move_optimizer_state(torch.device("cpu"))
[docs] @staticmethod def auto_device_count() -> int: """Get the devices when set to auto.""" return torch.cuda.device_count()
def get_nvidia_gpu_stats(device: torch.device) -> Dict[str, float]: """Get GPU stats including memory, fan speed, and temperature from nvidia-smi. Args: device: GPU device for which to get stats Returns: A dictionary mapping the metrics to their values. Raises: FileNotFoundError: If nvidia-smi installation not found """ nvidia_smi_path = shutil.which("nvidia-smi") if nvidia_smi_path is None: raise FileNotFoundError("nvidia-smi: command not found") gpu_stat_metrics = [ ("utilization.gpu", "%"), ("memory.used", "MB"), ("memory.free", "MB"), ("utilization.memory", "%"), ("fan.speed", "%"), ("temperature.gpu", "°C"), ("temperature.memory", "°C"), ] gpu_stat_keys = [k for k, _ in gpu_stat_metrics] gpu_query = ",".join(gpu_stat_keys) gpu_id = _get_gpu_id(device.index) result = subprocess.run( [nvidia_smi_path, f"--query-gpu={gpu_query}", "--format=csv,nounits,noheader", f"--id={gpu_id}"], encoding="utf-8", stdout=subprocess.PIPE, stderr=subprocess.PIPE, # for backward compatibility with python version 3.6 check=True, ) def _to_float(x: str) -> float: try: return float(x) except ValueError: return 0.0 s = result.stdout.strip() stats = [_to_float(x) for x in s.split(", ")] gpu_stats = {f"{x} ({unit})": stat for (x, unit), stat in zip(gpu_stat_metrics, stats)} return gpu_stats def _get_gpu_id(device_id: int) -> str: """Get the unmasked real GPU IDs.""" # All devices if `CUDA_VISIBLE_DEVICES` unset default = ",".join(str(i) for i in range(torch.cuda.device_count())) cuda_visible_devices: List[str] = os.getenv("CUDA_VISIBLE_DEVICES", default=default).split(",") return cuda_visible_devices[device_id].strip()

© Copyright Copyright (c) 2018-2021, William Falcon et al... Revision 46f718d2.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
1.5.4
1.5.3
1.5.2
1.5.1
1.5.0
1.4.9
1.4.8
1.4.7
1.4.6
1.4.5
1.4.4
1.4.3
1.4.2
1.4.1
1.4.0
1.3.8
1.3.7
1.3.6
1.3.5
1.3.4
1.3.3
1.3.2
1.3.1
1.3.0
1.2.10
1.2.8
1.2.7
1.2.6
1.2.5
1.2.4
1.2.3
1.2.2
1.2.1
1.2.0
1.1.8
1.1.7
1.1.6
1.1.5
1.1.4
1.1.3
1.1.2
1.1.1
1.1.0
1.0.8
1.0.7
1.0.6
1.0.5
1.0.4
1.0.3
1.0.2
1.0.1
1.0.0
0.10.0
0.9.0
0.8.5
0.8.4
0.8.3
0.8.2
0.8.1
0.8.0
0.7.6
0.7.5
0.7.4
0.7.3
0.7.2
0.7.1
0.7.0
0.6.0
0.5.3
0.4.9
ipynb-update
docs-search
Downloads
html
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.