File size: 793 Bytes
300c478 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
from __future__ import annotations
from numbers import Integral
import torch
def _validate_refinement_steps(value: object, *, source: str) -> int:
if isinstance(value, bool) or not isinstance(value, Integral):
raise TypeError(f'{source} must be a positive integer')
value = int(value)
if value < 1:
raise ValueError(f'{source} must be a positive integer')
return value
def refinement_mask_fraction(ages: torch.Tensor, *, steps: int) -> torch.Tensor:
steps = _validate_refinement_steps(steps, source='steps')
if ages.dtype == torch.bool or ages.is_floating_point():
raise TypeError('ages must use an integer dtype')
effective_age = ages.clamp(min=0, max=steps).to(dtype=torch.float32)
return (1.0 - effective_age / steps).mean(dim=-1)
|