Understanding PyTorch ceil() and floor() Methods

Rate this post

πŸ’‘ Problem Formulation: When working with numerical data in PyTorch, it’s common to need to round values up or down to the nearest integer. For instance, given a PyTorch tensor with floating-point values, you might want to obtain a new tensor where each element is the ceiling or floor of the original value to facilitate certain computations or to meet specific criteria. Ceil operations round up to the nearest integer, while floor operations round down.

Method 1: Using torch.ceil()

The torch.ceil() function in PyTorch returns a new tensor with the ceiling of the elements of the input tensor. The ceiling function rounds each element to the smallest integer greater than or equal to that number. This function is particularly useful when you want all your elements to round up in a tensor.

Here’s an example:

import torch

# Input tensor
input_tensor = torch.tensor([1.2, 3.7, -2.3])

# Apply the ceil function
output_tensor = torch.ceil(input_tensor)
print(output_tensor)

Output:

tensor([ 2.,  4., -2.])

In this code snippet, we use torch.ceil() to round up each element in our input tensor. The values 1.2 and 3.7 are rounded up to 2.0 and 4.0, respectively. Since -2.3 is already less than -2, it simply rounds up to -2.

Method 2: Using torch.floor()

The torch.floor() function performs the inverse operation of torch.ceil(). It returns a new tensor where each element is the largest integer less than or equal to the original value. Use this function when you need all elements in your tensor to round down.

Here’s an example:

import torch

# Input tensor
input_tensor = torch.tensor([1.2, 3.7, -2.3])

# Apply the floor function
output_tensor = torch.floor(input_tensor)
print(output_tensor)

Output:

tensor([ 1.,  3., -3.])

This snippet demonstrates the use of torch.floor() to round down the elements of the input tensor. The number 1.2 rounds down to 1.0, 3.7 to 3.0, and -2.3 to -3.0 because -3 is the largest integer less than -2.3.

Method 3: Manually implementing ceil and floor

Sometimes, for educational purposes or for customized behavior, you might want to implement your own ceil and floor operations. You can achieve this by utilizing basic arithmetic operations and the torch.trunc() function, which truncates each element in a tensor to its integer part.

Here’s an example:

import torch

# Input tensor
input_tensor = torch.tensor([1.2, 3.7, -2.3])

# Custom ceil by using trunc and adding 1
custom_ceil = torch.trunc(input_tensor) + ((input_tensor > 0).type(input_tensor.dtype))
print(custom_ceil)

# Custom floor by using trunc
custom_floor = torch.trunc(input_tensor) - ((input_tensor < 0).type(input_tensor.dtype))
print(custom_floor)

Output:

tensor([ 2.,  4., -2.])
tensor([ 1.,  3., -3.])

This example crafts a custom ceil and floor using torch.trunc(). For ceil, we add 1 to the truncated tensor where original values were positive, and for floor, we subtract 1 where they were negative. This mimics the behavior of regular ceil and floor functions with additional control if needed.

Method 4: In-place operations with ceil_() and floor_()

For optimization and memory efficiency, PyTorch provides in-place versions of ceil and floor operations: ceil_() and floor_(). These methods modify the input tensor directly, avoiding the creation of additional tensors.

Here’s an example:

import torch

# Input tensor
input_tensor = torch.tensor([1.2, 3.7, -2.3])

# Apply in-place ceil
input_tensor.ceil_()
print(input_tensor)

# Reset tensor and apply in-place floor
input_tensor.copy_(torch.tensor([1.2, 3.7, -2.3]))
input_tensor.floor_()
print(input_tensor)

Output:

tensor([ 2.,  4., -2.])
tensor([ 1.,  3., -3.])

Here, we first demonstrate an in-place ceil operation directly on the input tensor, and then we reset the tensor to its original values to show the in-place floor operation. These methods save memory but should be used with caution as they change the original tensor.

Bonus One-Liner Method 5: Using the round() function

Though not strictly ceiling or flooring, Python’s built-in round() function can be used for rounding tensors when the values are exactly halfway between two integers. This method rounds to the nearest even integer.

Here’s an example:

import torch

# Input tensor
input_tensor = torch.tensor([1.5, 2.5, 3.5])

# Apply round
rounded_tensor = torch.round(input_tensor)
print(rounded_tensor)

Output:

tensor([ 2.,  2.,  4.])

This snippet rounds the tensor elements using torch.round(). Interestingly, 1.5 and 3.5 round to 2 and 4 (nearest even numbers), while 2.5 also rounds to 2, demonstrating Python’s approach to tying breaking in rounding.

Summary/Discussion

  • Method 1: torch.ceil(). Straightforward method for rounding up. Clarity and ease of use. No control over in-place adjustments or custom rounding logic.
  • Method 2: torch.floor(). Direct counterpart to ceil, it rounds down. Similarly simple to ceil. Lacks nuanced control.
  • Method 3: Manual implementations. Offers flexibility for custom logic. More complex and error-prone than built-in methods.
  • Method 4: In-place operations. Memory efficient and fast. However, overwriting original data can be risky if it’s needed later.
  • Bonus Method 5: round(). Useful for standard rounding logic. May not be suitable for strict flooring or ceiling requirements.