实用工具

float_time_to_index(time, num_time_steps)

将浮点时间值转换为时间索引。

参数

名称	类型	描述	默认值
time	张量	范围在 [0, 1] 内的浮点时间值张量。	必需
num_time_steps	整数	离散时间步的数量。	必需

返回值

类型	描述
张量	torch.Tensor: 对应于输入浮点时间值的时间索引张量。

源代码位于 bionemo/moco/distributions/time/utils.py

def float_time_to_index(time: torch.Tensor, num_time_steps: int) -> torch.Tensor:
    """Convert a float time value to a time index.

    Args:
        time (torch.Tensor): A tensor of float time values in the range [0, 1].
        num_time_steps (int): The number of discrete time steps.

    Returns:
        torch.Tensor: A tensor of time indices corresponding to the input float time values.
    """
    # Ensure time values are in the range [0, 1]
    time = torch.clamp(time, 0.0, 1.0)

    # Scale to the index range and round
    indices = torch.round(time * (num_time_steps - 1)).to(torch.int64)

    return indices