PyTorch Where and Clamp
torch.where chooses values cell by cell, based on a condition. clamp limits every value to a range between a minimum and a maximum.
Choose with where
The condition is True for the two passing scores. where writes 1 there and 0 elsewhere.
import torch
scores = torch.tensor([45, 80, 62])
result = torch.where(scores >= 60, torch.tensor(1), torch.tensor(0))
print(result)Replace the negatives
where can also keep the original value on one side. Here the negatives become 0 and the rest stay.
values = torch.tensor([-5, 3, -1, 8])
print(torch.where(values < 0, torch.tensor(0), values))Cap a range with clamp
clamp needs no condition. Any value above the maximum comes down to the maximum.
scores = torch.tensor([45, 105, 88, 130])
print(scores.clamp(min=0, max=100))Example
import torch
scores = torch.tensor([-12, 45, 105, 88, -3, 130])
clean = torch.where(scores < 0, torch.tensor(0), scores)
capped = clean.clamp(max=100)
print("Raw scores:")
print(scores)
print("Negatives replaced with 0:")
print(clean)
print("Capped at 100:")
print(capped)Key points
- Condition: a bool tensor that says, cell by cell, which source to pick
- Use torch.where(condition, a, b) to take from a where the condition is True
- Use torch.where(condition, a, b) to take from b where the condition is False
- Use clamp(min=..., max=...) to pull values back inside a range
- A value inside the range passes through clamp unchanged
- In short, where swaps whole values and clamp trims values at the edges
In this example, scores mixes negative scores with scores above 100. torch.where finds each negative score and replaces it with 0. The other scores pass through where unchanged. clamp then caps the cleaned scores at 100. The three prints show the scores before, between, and after the two fixes.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
What does torch.where(values < 0, torch.tensor(0), values) do to a negative value?