YouTipPytorch Torch Device
## PyTorch `torch.device` Tutorial
In PyTorch, `torch.device` is an object that represents the device on which a `torch.Tensor` is allocated or will be allocated. It allows you to seamlessly write device-agnostic code that can run on both CPUs and GPUs (CUDA, MPS, etc.) without changing the underlying tensor operations.
---
## Syntax and Usage
The `torch.device` class can be instantiated using a string or a combination of a string and a device index.
### Constructor Signature
```python
torch.device(device)
# Or alternatively
torch.device(type, index=None)
```
### Parameters
* **`device`** *(string or torch.device)*: The device identifier.
* Common string formats include:
* `'cpu'`: Represents the CPU.
* `'cuda'`: Represents the default CUDA (GPU) device.
* `'cuda:0'`, `'cuda:1'`: Represents a specific CUDA device by its index.
* `'mps'`: Represents Apple Silicon GPU (Metal Performance Shaders) for macOS.
* **`type`** *(string)*: The device type (e.g., `'cpu'`, `'cuda'`, `'mps'`).
* **`index`** *(int, optional)*: The index of the device (e.g., `0` for the first GPU). If not specified, the current device for the device type is used.
---
## Code Examples
### 1. Basic Device Creation and Tensor Allocation
The following example demonstrates how to define CPU and CUDA devices, check for GPU availability, and allocate tensors directly to a specific device.
```python
import torch
# Create a CPU device
cpu_device = torch.device('cpu')
print("CPU Device:", cpu_device)
# Create a CUDA device if a GPU is available
if torch.cuda.is_available():
cuda_device = torch.device('cuda:0')
print("CUDA Device:", cuda_device)
# Create a tensor directly on the GPU
x = torch.zeros(3, 4, device=cuda_device)
print("Tensor Device:", x.device)
else:
print("CUDA is not available on this system.")
```
### 2. Writing Device-Agnostic Code
A best practice in PyTorch is to write code that automatically detects whether a GPU is available and falls back to the CPU if it is not.
```python
import torch
# Automatically select GPU (CUDA or MPS) if available, otherwise default to CPU
if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")
print(f"Using device: {device}")
# Create a tensor on the selected device
y = torch.ones((2, 2), device=device)
print(f"Tensor allocated on: {y.device}")
```
### 3. Moving Existing Tensors Between Devices
You can move an existing tensor to a target device using the `.to()` method.
```python
import torch
# Create a tensor on the CPU by default
x = torch.tensor([1.0, 2.0, 3.0])
print("Original device:", x.device)
# Move the tensor to GPU if available
if torch.cuda.is_available():
device = torch.device("cuda")
x_gpu = x.to(device)
print("New device:", x_gpu.device)
```
---
## Key Considerations
1. **Device Compatibility**: Operations between tensors (e.g., matrix multiplication, addition) can only be performed if both tensors reside on the **same device**. Attempting to perform operations on a CPU tensor and a GPU tensor will raise a `RuntimeError`.
2. **In-Place vs. Out-of-Place**: The `.to(device)` method is not in-place for tensors. It returns a new tensor copy on the specified device. However, for neural network modules (`torch.nn.Module`), calling `.to(device)` moves the model parameters in-place.
3. **Multi-GPU Systems**: On systems with multiple GPUs, you can target specific cards using `'cuda:0'`, `'cuda:1'`, etc. You can also use `torch.cuda.set_device(device_id)` to set the default GPU index.