Overview
Thetorch package contains data structures for multi-dimensional tensors and defines mathematical operations over these tensors. It provides utilities for efficient serialization and CUDA support for GPU computation.
Core Functions
Tensor Creation
torch.tensor()
torch.tensor()
Creates a tensor from data.
Initial data for the tensor. Can be a list, tuple, NumPy ndarray, scalar, and other types.
The desired data type of returned tensor. If None, infers data type from data.
The desired device of returned tensor. Default: CPU for CPU tensor types, current CUDA device for CUDA tensor types.
If autograd should record operations on the returned tensor.
If set, returned tensor would be allocated in pinned memory. Works only for CPU tensors.
torch.zeros()
torch.zeros()
Returns a tensor filled with the scalar value 0.
A sequence of integers defining the shape of the output tensor.
The desired data type. Default: global default dtype.
The desired device of returned tensor.
If autograd should record operations on the returned tensor.
torch.ones()
torch.ones()
torch.rand()
torch.rand()
Returns a tensor filled with random numbers from a uniform distribution on the interval [0, 1).
A sequence of integers defining the shape of the output tensor.
A pseudorandom number generator for sampling.
The desired data type. Default: global default dtype.
The desired device of returned tensor.
torch.randn()
torch.randn()
Returns a tensor filled with random numbers from a normal distribution with mean 0 and variance 1.
A sequence of integers defining the shape of the output tensor.
A pseudorandom number generator for sampling.
The desired data type. Default: global default dtype.
Mathematical Operations
torch.abs()
torch.abs()
torch.add()
torch.add()
torch.matmul()
torch.matmul()
Matrix product of two tensors with broadcasting support.The behavior depends on the dimensionality of the tensors:
The first tensor to be multiplied.
The second tensor to be multiplied.
The output tensor.
- If both tensors are 1-dimensional, the dot product (scalar) is returned.
- If both arguments are 2-dimensional, the matrix-matrix product is returned.
- If the first argument is 1-dimensional and the second argument is 2-dimensional, a 1 is prepended to its dimension for the purpose of the matrix multiply. After the matrix multiply, the prepended dimension is removed.
- If the first argument is 2-dimensional and the second argument is 1-dimensional, the matrix-vector product is returned.
torch.sum()
torch.sum()
Returns the sum of all elements in the input tensor.
The input tensor.
The dimension or dimensions to reduce. If None, all dimensions are reduced.
Whether the output tensor has dim retained or not.
The desired data type of returned tensor. If specified, the input tensor is casted to dtype before operation.
Utilities
torch.save()
torch.save()
torch.load()
torch.load()
Loads an object saved with torch.save() from a file.
A file-like object or a string containing a file name.
A function, torch.device, string or dict specifying how to remap storage locations.
If True, only loads tensors, primitive types, and dictionaries.
torch.manual_seed()
torch.manual_seed()
Sets the seed for generating random numbers. Returns a torch.Generator object.
The desired seed value.
torch.set_default_device()
torch.set_default_device()
Sets the default torch.Tensor device.
The device to set as default.
torch.set_float32_matmul_precision()
torch.set_float32_matmul_precision()
Sets the internal precision of float32 matrix multiplications.Running float32 matrix multiplications in lower precision may significantly increase performance:
Can be “highest”, “high”, or “medium”.
- “highest”: float32 datatype (24 mantissa bits) for internal computations
- “high”: TensorFloat32 datatype (10 mantissa bits) or bfloat16-based algorithm
- “medium”: bfloat16 datatype (8 mantissa bits) for internal computations
Configuration
torch.use_deterministic_algorithms()
torch.use_deterministic_algorithms()
Symbolic Types
SymInt, SymFloat, SymBool
PyTorch provides symbolic types for dynamic shape tracking and symbolic computation:Best Practices
Tensor Creation Performance
Tensor Creation Performance
- Use
torch.empty()when you’ll immediately fill the tensor - Use
torch.zeros_like()andtorch.ones_like()to match existing tensor properties - Specify
device='cuda'directly during creation to avoid CPU→GPU copy
Reproducibility
Reproducibility
Set all random seeds for reproducible results:
Memory Management
Memory Management
- Use
torch.no_grad()context when not training to reduce memory usage - Call
delon large tensors andtorch.cuda.empty_cache()to free GPU memory - Use gradient checkpointing for large models
Related APIs
Tensor API
Tensor class methods and operations
Autograd API
Automatic differentiation functions
CUDA API
GPU acceleration with CUDA