Skip to content

Inference time schedules

ContinuousInferenceSchedule

Bases: InferenceSchedule

A base class for continuous time inference schedules.

Source code in bionemo/moco/schedules/inference_time_schedules.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
class ContinuousInferenceSchedule(InferenceSchedule):
    """A base class for continuous time inference schedules."""

    def __init__(
        self,
        nsteps: int,
        inclusive_end: bool = False,
        min_t: Float = 0,
        padding: Float = 0,
        dilation: Float = 0,
        direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
        device: Union[str, torch.device] = "cpu",
    ):
        """Initialize the ContinuousInferenceSchedule.

        Args:
            nsteps (int): Number of time steps.
            inclusive_end (bool): If True, include the end value (1.0) in the schedule otherwise ends at 1.0-1/nsteps (default is False).
            min_t (Float): minimum time value defaults to 0.
            padding (Float): padding time value defaults to 0.
            dilation (Float): dilation time value defaults to 0 ie the number of replicates.
            direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
            device (Optional[str]): Device to place the schedule on (default is "cpu").

        """
        super().__init__(nsteps, min_t, padding, dilation, direction, device)
        self.inclusive_end = inclusive_end

    def discretize(
        self,
        nsteps: Optional[int] = None,
        schedule: Optional[Tensor] = None,
        device: Optional[Union[str, torch.device]] = None,
    ) -> Tensor:
        """Discretize the time schedule into a list of time deltas.

        Args:
            nsteps (Optioanl[int]): Number of time steps. If None, uses the value from initialization.
            schedule (Optional[Tensor]): Time scheudle if None will generate it with generate_schedule.
            device (Optional[str]): Device to place the schedule on (default is "cpu").

        Returns:
            Tensor: A tensor of time deltas.
        """
        if device is None:
            device = self.device
        if schedule is None:
            schedule = self.generate_schedule(nsteps, device=device)
        if self.direction == TimeDirection.UNIFIED:
            schedule = torch.cat((schedule, torch.ones((1,), device=schedule.device)))
            dt = schedule[1:] - schedule[:-1]
        else:
            schedule = torch.cat((schedule, torch.zeros((1,), device=schedule.device)))
            dt = -1 * (schedule[1:] - schedule[:-1])
        return dt

__init__(nsteps, inclusive_end=False, min_t=0, padding=0, dilation=0, direction=TimeDirection.UNIFIED, device='cpu')

Initialize the ContinuousInferenceSchedule.

Parameters:

Name Type Description Default
nsteps int

Number of time steps.

 required
inclusive_end bool

If True, include the end value (1.0) in the schedule otherwise ends at 1.0-1/nsteps (default is False).

 False
min_t Float

minimum time value defaults to 0.

 0
padding Float

padding time value defaults to 0.

 0
dilation Float

dilation time value defaults to 0 ie the number of replicates.

 0
direction Optional[str]

TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).

 UNIFIED
device Optional[str]

Device to place the schedule on (default is "cpu").

 'cpu'
Source code in bionemo/moco/schedules/inference_time_schedules.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def __init__(
    self,
    nsteps: int,
    inclusive_end: bool = False,
    min_t: Float = 0,
    padding: Float = 0,
    dilation: Float = 0,
    direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
    device: Union[str, torch.device] = "cpu",
):
    """Initialize the ContinuousInferenceSchedule.

    Args:
        nsteps (int): Number of time steps.
        inclusive_end (bool): If True, include the end value (1.0) in the schedule otherwise ends at 1.0-1/nsteps (default is False).
        min_t (Float): minimum time value defaults to 0.
        padding (Float): padding time value defaults to 0.
        dilation (Float): dilation time value defaults to 0 ie the number of replicates.
        direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
        device (Optional[str]): Device to place the schedule on (default is "cpu").

    """
    super().__init__(nsteps, min_t, padding, dilation, direction, device)
    self.inclusive_end = inclusive_end

discretize(nsteps=None, schedule=None, device=None)

Discretize the time schedule into a list of time deltas.

Parameters:

Name Type Description Default
nsteps Optioanl[int]

Number of time steps. If None, uses the value from initialization.

 None
schedule Optional[Tensor]

Time scheudle if None will generate it with generate_schedule.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 None

Returns:

Name Type Description
Tensor Tensor

A tensor of time deltas.
Source code in bionemo/moco/schedules/inference_time_schedules.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def discretize(
    self,
    nsteps: Optional[int] = None,
    schedule: Optional[Tensor] = None,
    device: Optional[Union[str, torch.device]] = None,
) -> Tensor:
    """Discretize the time schedule into a list of time deltas.

    Args:
        nsteps (Optioanl[int]): Number of time steps. If None, uses the value from initialization.
        schedule (Optional[Tensor]): Time scheudle if None will generate it with generate_schedule.
        device (Optional[str]): Device to place the schedule on (default is "cpu").

    Returns:
        Tensor: A tensor of time deltas.
    """
    if device is None:
        device = self.device
    if schedule is None:
        schedule = self.generate_schedule(nsteps, device=device)
    if self.direction == TimeDirection.UNIFIED:
        schedule = torch.cat((schedule, torch.ones((1,), device=schedule.device)))
        dt = schedule[1:] - schedule[:-1]
    else:
        schedule = torch.cat((schedule, torch.zeros((1,), device=schedule.device)))
        dt = -1 * (schedule[1:] - schedule[:-1])
    return dt

DiscreteInferenceSchedule

Bases: InferenceSchedule

A base class for discrete time inference schedules.

Source code in bionemo/moco/schedules/inference_time_schedules.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
class DiscreteInferenceSchedule(InferenceSchedule):
    """A base class for discrete time inference schedules."""

    def discretize(
        self,
        nsteps: Optional[int] = None,
        device: Optional[Union[str, torch.device]] = None,
    ) -> Tensor:
        """Discretize the time schedule into a list of time deltas.

        Args:
            nsteps (Optioanl[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").

        Returns:
            Tensor: A tensor of time deltas.
        """
        if self.padding > 0 or self.dilation > 0:
            raise NotImplementedError("discreteize is not implemented for discrete schedules with padding or dilation")
        if device is None:
            device = self.device
        return torch.full(
            (nsteps if nsteps is not None else self.nsteps,),
            1 / (nsteps if nsteps is not None else self.nsteps),
            device=device,
        )

discretize(nsteps=None, device=None)

Discretize the time schedule into a list of time deltas.

Parameters:

Name Type Description Default
nsteps Optioanl[int]

Number of time steps. If None, uses the value from initialization.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 None

Returns:

Name Type Description
Tensor Tensor

A tensor of time deltas.
Source code in bionemo/moco/schedules/inference_time_schedules.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def discretize(
    self,
    nsteps: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
) -> Tensor:
    """Discretize the time schedule into a list of time deltas.

    Args:
        nsteps (Optioanl[int]): Number of time steps. If None, uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").

    Returns:
        Tensor: A tensor of time deltas.
    """
    if self.padding > 0 or self.dilation > 0:
        raise NotImplementedError("discreteize is not implemented for discrete schedules with padding or dilation")
    if device is None:
        device = self.device
    return torch.full(
        (nsteps if nsteps is not None else self.nsteps,),
        1 / (nsteps if nsteps is not None else self.nsteps),
        device=device,
    )

DiscreteLinearInferenceSchedule

Bases: DiscreteInferenceSchedule

A linear time schedule for discrete time inference.

Source code in bionemo/moco/schedules/inference_time_schedules.py 
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
class DiscreteLinearInferenceSchedule(DiscreteInferenceSchedule):
    """A linear time schedule for discrete time inference."""

    def __init__(
        self,
        nsteps: int,
        min_t: Float = 0,
        padding: Float = 0,
        dilation: Float = 0,
        direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
        device: Union[str, torch.device] = "cpu",
    ):
        """Initialize the DiscreteLinearInferenceSchedule.

        Args:
            nsteps (int): Number of time steps.
            min_t (Float): minimum time value defaults to 0.
            padding (Float): padding time value defaults to 0.
            dilation (Float): dilation time value defaults to 0 ie the number of replicates.
            direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        super().__init__(nsteps, min_t, padding, dilation, direction, device)

    def generate_schedule(
        self, nsteps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None
    ) -> Tensor:
        """Generate the linear time schedule as a tensor.

        Args:
            nsteps (Optional[int]): Number of time steps. If None uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").

        Returns:
            Tensor: A tensor of time steps.
            Tensor: A tensor of time steps.
        """
        if device is None:
            device = self.device
        if nsteps is None:
            nsteps = self.nsteps
        nsteps -= self.padding
        dilation = self.dilation + 1
        if dilation > 1:
            if nsteps % dilation != 0:
                raise ValueError(f"nsteps ({nsteps}) is not divisible by dilation + 1 ({dilation})")
            nsteps = int(nsteps / self.dilation)
        if nsteps is None:
            raise ValueError("nsteps cannot be None")
        schedule = torch.arange(nsteps).to(device=device)
        if dilation > 1:
            schedule = schedule.repeat_interleave(dilation)
        if self.direction == TimeDirection.DIFFUSION:
            schedule = schedule.flip(0)
        if self.padding > 0:
            schedule = torch.cat((schedule, schedule[-1] * torch.ones(self.padding, device=device)))
        return schedule

__init__(nsteps, min_t=0, padding=0, dilation=0, direction=TimeDirection.UNIFIED, device='cpu')

Initialize the DiscreteLinearInferenceSchedule.

Parameters:

Name Type Description Default
nsteps int

Number of time steps.

 required
min_t Float

minimum time value defaults to 0.

 0
padding Float

padding time value defaults to 0.

 0
dilation Float

dilation time value defaults to 0 ie the number of replicates.

 0
direction Optional[str]

TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).

 UNIFIED
device Optional[str]

Device to place the schedule on (default is "cpu").

 'cpu'
Source code in bionemo/moco/schedules/inference_time_schedules.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def __init__(
    self,
    nsteps: int,
    min_t: Float = 0,
    padding: Float = 0,
    dilation: Float = 0,
    direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
    device: Union[str, torch.device] = "cpu",
):
    """Initialize the DiscreteLinearInferenceSchedule.

    Args:
        nsteps (int): Number of time steps.
        min_t (Float): minimum time value defaults to 0.
        padding (Float): padding time value defaults to 0.
        dilation (Float): dilation time value defaults to 0 ie the number of replicates.
        direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
        device (Optional[str]): Device to place the schedule on (default is "cpu").
    """
    super().__init__(nsteps, min_t, padding, dilation, direction, device)

generate_schedule(nsteps=None, device=None)

Generate the linear time schedule as a tensor.

Parameters:

Name Type Description Default
nsteps Optional[int]

Number of time steps. If None uses the value from initialization.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 None

Returns:

Name Type Description
Tensor Tensor

A tensor of time steps.
Tensor Tensor

A tensor of time steps.
Source code in bionemo/moco/schedules/inference_time_schedules.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def generate_schedule(
    self, nsteps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None
) -> Tensor:
    """Generate the linear time schedule as a tensor.

    Args:
        nsteps (Optional[int]): Number of time steps. If None uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").

    Returns:
        Tensor: A tensor of time steps.
        Tensor: A tensor of time steps.
    """
    if device is None:
        device = self.device
    if nsteps is None:
        nsteps = self.nsteps
    nsteps -= self.padding
    dilation = self.dilation + 1
    if dilation > 1:
        if nsteps % dilation != 0:
            raise ValueError(f"nsteps ({nsteps}) is not divisible by dilation + 1 ({dilation})")
        nsteps = int(nsteps / self.dilation)
    if nsteps is None:
        raise ValueError("nsteps cannot be None")
    schedule = torch.arange(nsteps).to(device=device)
    if dilation > 1:
        schedule = schedule.repeat_interleave(dilation)
    if self.direction == TimeDirection.DIFFUSION:
        schedule = schedule.flip(0)
    if self.padding > 0:
        schedule = torch.cat((schedule, schedule[-1] * torch.ones(self.padding, device=device)))
    return schedule

InferenceSchedule

Bases: ABC

A base class for inference time schedules.

Source code in bionemo/moco/schedules/inference_time_schedules.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class InferenceSchedule(ABC):
    """A base class for inference time schedules."""

    def __init__(
        self,
        nsteps: int,
        min_t: Float = 0,
        padding: Float = 0,
        dilation: Float = 0,
        direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
        device: Union[str, torch.device] = "cpu",
    ):
        """Initialize the InferenceSchedule.

        Args:
            nsteps (int): Number of time steps.
            min_t (Float): minimum time value defaults to 0.
            padding (Float): padding time value defaults to 0.
            dilation (Float): dilation time value defaults to 0 ie the number of replicates.
            direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
            device (Optional[str]): Device to place the schedule on (default is "cpu").

        """
        self.nsteps = nsteps
        self.min_t = min_t
        self.padding = padding
        self.dilation = dilation
        self.direction = string_to_enum(direction, TimeDirection)
        self.device = device

    @abstractmethod
    def generate_schedule(
        self, nsteps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None
    ) -> Tensor:
        """Generate the time schedule as a tensor.

        Args:
            nsteps (Optioanl[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        pass

    def pad_time(
        self, n_samples: int, scalar_time: Float, device: Optional[Union[str, torch.device]] = None
    ) -> Tensor:
        """Creates a tensor of shape (n_samples,) filled with a scalar time value.

        Args:
            n_samples (int): The desired dimension of the output tensor.
            scalar_time (Float): The scalar time value to fill the tensor with.
            device (Optional[Union[str, torch.device]], optional):
                The device to place the tensor on. Defaults to None, which uses the default device.

        Returns:
            Tensor: A tensor of shape (n_samples,) filled with the scalar time value.
        """
        return torch.full((n_samples,), fill_value=scalar_time).to(device)

__init__(nsteps, min_t=0, padding=0, dilation=0, direction=TimeDirection.UNIFIED, device='cpu')

Initialize the InferenceSchedule.

Parameters:

Name Type Description Default
nsteps int

Number of time steps.

 required
min_t Float

minimum time value defaults to 0.

 0
padding Float

padding time value defaults to 0.

 0
dilation Float

dilation time value defaults to 0 ie the number of replicates.

 0
direction Optional[str]

TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).

 UNIFIED
device Optional[str]

Device to place the schedule on (default is "cpu").

 'cpu'
Source code in bionemo/moco/schedules/inference_time_schedules.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def __init__(
    self,
    nsteps: int,
    min_t: Float = 0,
    padding: Float = 0,
    dilation: Float = 0,
    direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
    device: Union[str, torch.device] = "cpu",
):
    """Initialize the InferenceSchedule.

    Args:
        nsteps (int): Number of time steps.
        min_t (Float): minimum time value defaults to 0.
        padding (Float): padding time value defaults to 0.
        dilation (Float): dilation time value defaults to 0 ie the number of replicates.
        direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
        device (Optional[str]): Device to place the schedule on (default is "cpu").

    """
    self.nsteps = nsteps
    self.min_t = min_t
    self.padding = padding
    self.dilation = dilation
    self.direction = string_to_enum(direction, TimeDirection)
    self.device = device

generate_schedule(nsteps=None, device=None) abstractmethod

Generate the time schedule as a tensor.

Parameters:

Name Type Description Default
nsteps Optioanl[int]

Number of time steps. If None, uses the value from initialization.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 None
Source code in bionemo/moco/schedules/inference_time_schedules.py 
58
59
60
61
62
63
64
65
66
67
68
@abstractmethod
def generate_schedule(
    self, nsteps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None
) -> Tensor:
    """Generate the time schedule as a tensor.

    Args:
        nsteps (Optioanl[int]): Number of time steps. If None, uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").
    """
    pass

pad_time(n_samples, scalar_time, device=None)

Creates a tensor of shape (n_samples,) filled with a scalar time value.

Parameters:

Name Type Description Default
n_samples int

The desired dimension of the output tensor.

 required
scalar_time Float

The scalar time value to fill the tensor with.

 required
device Optional[Union[str, device]]

The device to place the tensor on. Defaults to None, which uses the default device.

 None

Returns:

Name Type Description
Tensor Tensor

A tensor of shape (n_samples,) filled with the scalar time value.
Source code in bionemo/moco/schedules/inference_time_schedules.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def pad_time(
    self, n_samples: int, scalar_time: Float, device: Optional[Union[str, torch.device]] = None
) -> Tensor:
    """Creates a tensor of shape (n_samples,) filled with a scalar time value.

    Args:
        n_samples (int): The desired dimension of the output tensor.
        scalar_time (Float): The scalar time value to fill the tensor with.
        device (Optional[Union[str, torch.device]], optional):
            The device to place the tensor on. Defaults to None, which uses the default device.

    Returns:
        Tensor: A tensor of shape (n_samples,) filled with the scalar time value.
    """
    return torch.full((n_samples,), fill_value=scalar_time).to(device)

LinearInferenceSchedule

Bases: ContinuousInferenceSchedule

A linear time schedule for continuous time inference.

Source code in bionemo/moco/schedules/inference_time_schedules.py 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
class LinearInferenceSchedule(ContinuousInferenceSchedule):
    """A linear time schedule for continuous time inference."""

    def __init__(
        self,
        nsteps: int,
        inclusive_end: bool = False,
        min_t: Float = 0,
        padding: Float = 0,
        dilation: Float = 0,
        direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
        device: Union[str, torch.device] = "cpu",
    ):
        """Initialize the LinearInferenceSchedule.

        Args:
            nsteps (int): Number of time steps.
            inclusive_end (bool): If True, include the end value (1.0) in the schedule otherwise ends at 1.0-1/nsteps (default is False).
            min_t (Float): minimum time value defaults to 0.
            padding (Float): padding time value defaults to 0.
            dilation (Float): dilation time value defaults to 0 ie the number of replicates.
            direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        super().__init__(nsteps, inclusive_end, min_t, padding, dilation, direction, device)

    def generate_schedule(
        self,
        nsteps: Optional[int] = None,
        device: Optional[Union[str, torch.device]] = None,
    ) -> Tensor:
        """Generate the linear time schedule as a tensor.

        Args:
            nsteps (Optional[int]): Number of time steps. If None uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").

        Returns:
            Tensor: A tensor of time steps.
        """
        if device is None:
            device = self.device
        if nsteps is None:
            nsteps = self.nsteps
        nsteps -= self.padding
        dilation = self.dilation + 1
        if dilation > 1:
            if nsteps % dilation != 0:
                raise ValueError(f"nsteps ({nsteps}) is not divisible by dilation + 1 ({dilation})")
            nsteps = int(nsteps / dilation)
        if nsteps is None:
            raise ValueError("nsteps cannot be None")
        if not self.inclusive_end:
            schedule = torch.linspace(self.min_t, 1, nsteps + 1).to(device=device)
            schedule = schedule[:-1]
        else:
            schedule = torch.linspace(self.min_t, 1, nsteps).to(device=device)
        if dilation > 1:
            schedule = schedule.repeat_interleave(dilation)
        if self.padding > 0:
            schedule = torch.cat((schedule, torch.ones(self.padding, device=device)))
        if self.direction == TimeDirection.DIFFUSION:
            schedule = 1 - schedule
        return schedule

__init__(nsteps, inclusive_end=False, min_t=0, padding=0, dilation=0, direction=TimeDirection.UNIFIED, device='cpu')

Initialize the LinearInferenceSchedule.

Parameters:

Name Type Description Default
nsteps int

Number of time steps.

 required
inclusive_end bool

If True, include the end value (1.0) in the schedule otherwise ends at 1.0-1/nsteps (default is False).

 False
min_t Float

minimum time value defaults to 0.

 0
padding Float

padding time value defaults to 0.

 0
dilation Float

dilation time value defaults to 0 ie the number of replicates.

 0
direction Optional[str]

TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).

 UNIFIED
device Optional[str]

Device to place the schedule on (default is "cpu").

 'cpu'
Source code in bionemo/moco/schedules/inference_time_schedules.py 
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
def __init__(
    self,
    nsteps: int,
    inclusive_end: bool = False,
    min_t: Float = 0,
    padding: Float = 0,
    dilation: Float = 0,
    direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
    device: Union[str, torch.device] = "cpu",
):
    """Initialize the LinearInferenceSchedule.

    Args:
        nsteps (int): Number of time steps.
        inclusive_end (bool): If True, include the end value (1.0) in the schedule otherwise ends at 1.0-1/nsteps (default is False).
        min_t (Float): minimum time value defaults to 0.
        padding (Float): padding time value defaults to 0.
        dilation (Float): dilation time value defaults to 0 ie the number of replicates.
        direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
        device (Optional[str]): Device to place the schedule on (default is "cpu").
    """
    super().__init__(nsteps, inclusive_end, min_t, padding, dilation, direction, device)

generate_schedule(nsteps=None, device=None)

Generate the linear time schedule as a tensor.

Parameters:

Name Type Description Default
nsteps Optional[int]

Number of time steps. If None uses the value from initialization.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 None

Returns:

Name Type Description
Tensor Tensor

A tensor of time steps.
Source code in bionemo/moco/schedules/inference_time_schedules.py 
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
def generate_schedule(
    self,
    nsteps: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
) -> Tensor:
    """Generate the linear time schedule as a tensor.

    Args:
        nsteps (Optional[int]): Number of time steps. If None uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").

    Returns:
        Tensor: A tensor of time steps.
    """
    if device is None:
        device = self.device
    if nsteps is None:
        nsteps = self.nsteps
    nsteps -= self.padding
    dilation = self.dilation + 1
    if dilation > 1:
        if nsteps % dilation != 0:
            raise ValueError(f"nsteps ({nsteps}) is not divisible by dilation + 1 ({dilation})")
        nsteps = int(nsteps / dilation)
    if nsteps is None:
        raise ValueError("nsteps cannot be None")
    if not self.inclusive_end:
        schedule = torch.linspace(self.min_t, 1, nsteps + 1).to(device=device)
        schedule = schedule[:-1]
    else:
        schedule = torch.linspace(self.min_t, 1, nsteps).to(device=device)
    if dilation > 1:
        schedule = schedule.repeat_interleave(dilation)
    if self.padding > 0:
        schedule = torch.cat((schedule, torch.ones(self.padding, device=device)))
    if self.direction == TimeDirection.DIFFUSION:
        schedule = 1 - schedule
    return schedule

LogInferenceSchedule

Bases: ContinuousInferenceSchedule

A log time schedule for inference, where time steps are generated by taking the logarithm of a uniform schedule.

Source code in bionemo/moco/schedules/inference_time_schedules.py 
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
class LogInferenceSchedule(ContinuousInferenceSchedule):
    """A log time schedule for inference, where time steps are generated by taking the logarithm of a uniform schedule."""

    def __init__(
        self,
        nsteps: int,
        inclusive_end: bool = False,
        min_t: Float = 0,
        padding: Float = 0,
        dilation: Float = 0,
        exponent: Float = -2.0,
        direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
        device: Union[str, torch.device] = "cpu",
    ):
        """Initialize the LogInferenceSchedule.

        Returns a log space time schedule.

        Which for 100 steps with default parameters is:
            tensor([0.0000, 0.0455, 0.0889, 0.1303, 0.1699, 0.2077, 0.2439, 0.2783, 0.3113,
                    0.3427, 0.3728, 0.4015, 0.4288, 0.4550, 0.4800, 0.5039, 0.5266, 0.5484,
                    0.5692, 0.5890, 0.6080, 0.6261, 0.6434, 0.6599, 0.6756, 0.6907, 0.7051,
                    0.7188, 0.7319, 0.7444, 0.7564, 0.7678, 0.7787, 0.7891, 0.7991, 0.8086,
                    0.8176, 0.8263, 0.8346, 0.8425, 0.8500, 0.8572, 0.8641, 0.8707, 0.8769,
                    0.8829, 0.8887, 0.8941, 0.8993, 0.9043, 0.9091, 0.9136, 0.9180, 0.9221,
                    0.9261, 0.9299, 0.9335, 0.9369, 0.9402, 0.9434, 0.9464, 0.9492, 0.9520,
                    0.9546, 0.9571, 0.9595, 0.9618, 0.9639, 0.9660, 0.9680, 0.9699, 0.9717,
                    0.9734, 0.9751, 0.9767, 0.9782, 0.9796, 0.9810, 0.9823, 0.9835, 0.9847,
                    0.9859, 0.9870, 0.9880, 0.9890, 0.9899, 0.9909, 0.9917, 0.9925, 0.9933,
                    0.9941, 0.9948, 0.9955, 0.9962, 0.9968, 0.9974, 0.9980, 0.9985, 0.9990,
                    0.9995])

        Args:
            nsteps (int): Number of time steps.
            inclusive_end (bool): If True, include the end value (1.0) in the schedule otherwise ends at <1.0 (default is False).
            min_t (Float): minimum time value defaults to 0.
            padding (Float): padding time value defaults to 0.
            dilation (Float): dilation time value defaults to 0 ie the number of replicates.
            exponent (Float): log space exponent parameter defaults to -2.0. The lower number the more aggressive the acceleration of 0 to 0.9 will be thus having more steps from 0.9 to 1.0.
            direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        super().__init__(nsteps, inclusive_end, min_t, padding, dilation, direction, device)
        if exponent is None:
            raise ValueError("exponent cannot be None for the log schedule")
        if exponent >= 0:
            raise ValueError(f"exponent input must be <0, got {exponent}")
        self.exponent = exponent

    def generate_schedule(
        self,
        nsteps: Optional[int] = None,
        device: Optional[Union[str, torch.device]] = None,
    ) -> Tensor:
        """Generate the log time schedule as a tensor.

        Args:
            nsteps (Optional[int]): Number of time steps. If None uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        if device is None:
            device = self.device
        if nsteps is None:
            nsteps = self.nsteps
        nsteps -= self.padding
        dilation = self.dilation + 1
        if dilation > 1:
            if nsteps % dilation != 0:
                raise ValueError(f"nsteps ({nsteps}) is not divisible by dilation + 1 ({dilation})")
            nsteps = int(nsteps / self.dilation)
        if nsteps is None:
            raise ValueError("nsteps cannot be None")

        if not self.inclusive_end:
            t = 1.0 - torch.logspace(self.exponent, 0, nsteps + 1).flip(0).to(device=device)
            t = t - torch.min(t)
            schedule = t / torch.max(t)
            schedule = schedule[:-1]
        else:
            t = 1.0 - torch.logspace(self.exponent, 0, nsteps).flip(0).to(device=device)
            t = t - torch.min(t)
            schedule = t / torch.max(t)

        if self.min_t > 0:
            schedule = torch.clamp(schedule, min=self.min_t)

        if dilation > 1:
            schedule = schedule.repeat_interleave(dilation)
        if self.padding > 0:
            schedule = torch.cat((schedule, torch.ones(self.padding, device=device)))
        if self.direction == TimeDirection.DIFFUSION:
            schedule = 1 - schedule
        return schedule

__init__(nsteps, inclusive_end=False, min_t=0, padding=0, dilation=0, exponent=-2.0, direction=TimeDirection.UNIFIED, device='cpu')

Initialize the LogInferenceSchedule.

Returns a log space time schedule.

Which for 100 steps with default parameters is

tensor([0.0000, 0.0455, 0.0889, 0.1303, 0.1699, 0.2077, 0.2439, 0.2783, 0.3113, 0.3427, 0.3728, 0.4015, 0.4288, 0.4550, 0.4800, 0.5039, 0.5266, 0.5484, 0.5692, 0.5890, 0.6080, 0.6261, 0.6434, 0.6599, 0.6756, 0.6907, 0.7051, 0.7188, 0.7319, 0.7444, 0.7564, 0.7678, 0.7787, 0.7891, 0.7991, 0.8086, 0.8176, 0.8263, 0.8346, 0.8425, 0.8500, 0.8572, 0.8641, 0.8707, 0.8769, 0.8829, 0.8887, 0.8941, 0.8993, 0.9043, 0.9091, 0.9136, 0.9180, 0.9221, 0.9261, 0.9299, 0.9335, 0.9369, 0.9402, 0.9434, 0.9464, 0.9492, 0.9520, 0.9546, 0.9571, 0.9595, 0.9618, 0.9639, 0.9660, 0.9680, 0.9699, 0.9717, 0.9734, 0.9751, 0.9767, 0.9782, 0.9796, 0.9810, 0.9823, 0.9835, 0.9847, 0.9859, 0.9870, 0.9880, 0.9890, 0.9899, 0.9909, 0.9917, 0.9925, 0.9933, 0.9941, 0.9948, 0.9955, 0.9962, 0.9968, 0.9974, 0.9980, 0.9985, 0.9990, 0.9995])

Parameters:

Name Type Description Default
nsteps int

Number of time steps.

 required
inclusive_end bool

If True, include the end value (1.0) in the schedule otherwise ends at <1.0 (default is False).

 False
min_t Float

minimum time value defaults to 0.

 0
padding Float

padding time value defaults to 0.

 0
dilation Float

dilation time value defaults to 0 ie the number of replicates.

 0
exponent Float

log space exponent parameter defaults to -2.0. The lower number the more aggressive the acceleration of 0 to 0.9 will be thus having more steps from 0.9 to 1.0.

 -2.0
direction Optional[str]

TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).

 UNIFIED
device Optional[str]

Device to place the schedule on (default is "cpu").

 'cpu'
Source code in bionemo/moco/schedules/inference_time_schedules.py 
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
def __init__(
    self,
    nsteps: int,
    inclusive_end: bool = False,
    min_t: Float = 0,
    padding: Float = 0,
    dilation: Float = 0,
    exponent: Float = -2.0,
    direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
    device: Union[str, torch.device] = "cpu",
):
    """Initialize the LogInferenceSchedule.

    Returns a log space time schedule.

    Which for 100 steps with default parameters is:
        tensor([0.0000, 0.0455, 0.0889, 0.1303, 0.1699, 0.2077, 0.2439, 0.2783, 0.3113,
                0.3427, 0.3728, 0.4015, 0.4288, 0.4550, 0.4800, 0.5039, 0.5266, 0.5484,
                0.5692, 0.5890, 0.6080, 0.6261, 0.6434, 0.6599, 0.6756, 0.6907, 0.7051,
                0.7188, 0.7319, 0.7444, 0.7564, 0.7678, 0.7787, 0.7891, 0.7991, 0.8086,
                0.8176, 0.8263, 0.8346, 0.8425, 0.8500, 0.8572, 0.8641, 0.8707, 0.8769,
                0.8829, 0.8887, 0.8941, 0.8993, 0.9043, 0.9091, 0.9136, 0.9180, 0.9221,
                0.9261, 0.9299, 0.9335, 0.9369, 0.9402, 0.9434, 0.9464, 0.9492, 0.9520,
                0.9546, 0.9571, 0.9595, 0.9618, 0.9639, 0.9660, 0.9680, 0.9699, 0.9717,
                0.9734, 0.9751, 0.9767, 0.9782, 0.9796, 0.9810, 0.9823, 0.9835, 0.9847,
                0.9859, 0.9870, 0.9880, 0.9890, 0.9899, 0.9909, 0.9917, 0.9925, 0.9933,
                0.9941, 0.9948, 0.9955, 0.9962, 0.9968, 0.9974, 0.9980, 0.9985, 0.9990,
                0.9995])

    Args:
        nsteps (int): Number of time steps.
        inclusive_end (bool): If True, include the end value (1.0) in the schedule otherwise ends at <1.0 (default is False).
        min_t (Float): minimum time value defaults to 0.
        padding (Float): padding time value defaults to 0.
        dilation (Float): dilation time value defaults to 0 ie the number of replicates.
        exponent (Float): log space exponent parameter defaults to -2.0. The lower number the more aggressive the acceleration of 0 to 0.9 will be thus having more steps from 0.9 to 1.0.
        direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
        device (Optional[str]): Device to place the schedule on (default is "cpu").
    """
    super().__init__(nsteps, inclusive_end, min_t, padding, dilation, direction, device)
    if exponent is None:
        raise ValueError("exponent cannot be None for the log schedule")
    if exponent >= 0:
        raise ValueError(f"exponent input must be <0, got {exponent}")
    self.exponent = exponent

generate_schedule(nsteps=None, device=None)

Generate the log time schedule as a tensor.

Parameters:

Name Type Description Default
nsteps Optional[int]

Number of time steps. If None uses the value from initialization.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 None
Source code in bionemo/moco/schedules/inference_time_schedules.py 
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
def generate_schedule(
    self,
    nsteps: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
) -> Tensor:
    """Generate the log time schedule as a tensor.

    Args:
        nsteps (Optional[int]): Number of time steps. If None uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").
    """
    if device is None:
        device = self.device
    if nsteps is None:
        nsteps = self.nsteps
    nsteps -= self.padding
    dilation = self.dilation + 1
    if dilation > 1:
        if nsteps % dilation != 0:
            raise ValueError(f"nsteps ({nsteps}) is not divisible by dilation + 1 ({dilation})")
        nsteps = int(nsteps / self.dilation)
    if nsteps is None:
        raise ValueError("nsteps cannot be None")

    if not self.inclusive_end:
        t = 1.0 - torch.logspace(self.exponent, 0, nsteps + 1).flip(0).to(device=device)
        t = t - torch.min(t)
        schedule = t / torch.max(t)
        schedule = schedule[:-1]
    else:
        t = 1.0 - torch.logspace(self.exponent, 0, nsteps).flip(0).to(device=device)
        t = t - torch.min(t)
        schedule = t / torch.max(t)

    if self.min_t > 0:
        schedule = torch.clamp(schedule, min=self.min_t)

    if dilation > 1:
        schedule = schedule.repeat_interleave(dilation)
    if self.padding > 0:
        schedule = torch.cat((schedule, torch.ones(self.padding, device=device)))
    if self.direction == TimeDirection.DIFFUSION:
        schedule = 1 - schedule
    return schedule

PowerInferenceSchedule

Bases: ContinuousInferenceSchedule

A power time schedule for inference, where time steps are generated by raising a uniform schedule to a specified power.

Source code in bionemo/moco/schedules/inference_time_schedules.py 
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
class PowerInferenceSchedule(ContinuousInferenceSchedule):
    """A power time schedule for inference, where time steps are generated by raising a uniform schedule to a specified power."""

    def __init__(
        self,
        nsteps: int,
        inclusive_end: bool = False,
        min_t: Float = 0,
        padding: Float = 0,
        dilation: Float = 0,
        exponent: Float = 1.0,
        direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
        device: Union[str, torch.device] = "cpu",
    ):
        """Initialize the PowerInferenceSchedule.

        Args:
            nsteps (int): Number of time steps.
            inclusive_end (bool): If True, include the end value (1.0) in the schedule otherwise ends at <1.0 (default is False).
            min_t (Float): minimum time value defaults to 0.
            padding (Float): padding time value defaults to 0.
            dilation (Float): dilation time value defaults to 0 ie the number of replicates.
            exponent (Float): Power parameter defaults to 1.0.
            direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        super().__init__(nsteps, inclusive_end, min_t, padding, dilation, direction, device)
        self.exponent = exponent

    def generate_schedule(
        self,
        nsteps: Optional[int] = None,
        device: Optional[Union[str, torch.device]] = None,
    ) -> Tensor:
        """Generate the power time schedule as a tensor.

        Args:
            nsteps (Optional[int]): Number of time steps. If None uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").


        Returns:
            Tensor: A tensor of time steps.
            Tensor: A tensor of time steps.
        """
        if device is None:
            device = self.device
        if nsteps is None:
            nsteps = self.nsteps
        nsteps -= self.padding
        dilation = self.dilation + 1
        if dilation > 1:
            if nsteps % dilation != 0:
                raise ValueError(f"nsteps ({nsteps}) is not divisible by dilation + 1 ({dilation})")
            nsteps = int(nsteps / dilation)
        if nsteps is None:
            raise ValueError("nsteps cannot be None")
        if not self.inclusive_end:
            schedule = torch.linspace(self.min_t, 1, nsteps + 1).to(device=device) ** self.exponent
            schedule = schedule[:-1]
        else:
            schedule = torch.linspace(self.min_t, 1, nsteps).to(device=device) ** self.exponent
        if dilation > 1:
            schedule = schedule.repeat_interleave(dilation)
        if self.padding > 0:
            schedule = torch.cat((schedule, torch.ones(self.padding, device=device)))
        if self.direction == TimeDirection.DIFFUSION:
            schedule = 1 - schedule
        return schedule

__init__(nsteps, inclusive_end=False, min_t=0, padding=0, dilation=0, exponent=1.0, direction=TimeDirection.UNIFIED, device='cpu')

Initialize the PowerInferenceSchedule.

Parameters:

Name Type Description Default
nsteps int

Number of time steps.

 required
inclusive_end bool

If True, include the end value (1.0) in the schedule otherwise ends at <1.0 (default is False).

 False
min_t Float

minimum time value defaults to 0.

 0
padding Float

padding time value defaults to 0.

 0
dilation Float

dilation time value defaults to 0 ie the number of replicates.

 0
exponent Float

Power parameter defaults to 1.0.

 1.0
direction Optional[str]

TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).

 UNIFIED
device Optional[str]

Device to place the schedule on (default is "cpu").

 'cpu'
Source code in bionemo/moco/schedules/inference_time_schedules.py 
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
def __init__(
    self,
    nsteps: int,
    inclusive_end: bool = False,
    min_t: Float = 0,
    padding: Float = 0,
    dilation: Float = 0,
    exponent: Float = 1.0,
    direction: Union[TimeDirection, str] = TimeDirection.UNIFIED,
    device: Union[str, torch.device] = "cpu",
):
    """Initialize the PowerInferenceSchedule.

    Args:
        nsteps (int): Number of time steps.
        inclusive_end (bool): If True, include the end value (1.0) in the schedule otherwise ends at <1.0 (default is False).
        min_t (Float): minimum time value defaults to 0.
        padding (Float): padding time value defaults to 0.
        dilation (Float): dilation time value defaults to 0 ie the number of replicates.
        exponent (Float): Power parameter defaults to 1.0.
        direction (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one (default is None).
        device (Optional[str]): Device to place the schedule on (default is "cpu").
    """
    super().__init__(nsteps, inclusive_end, min_t, padding, dilation, direction, device)
    self.exponent = exponent

generate_schedule(nsteps=None, device=None)

Generate the power time schedule as a tensor.

Parameters:

Name Type Description Default
nsteps Optional[int]

Number of time steps. If None uses the value from initialization.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 None

Returns:

Name Type Description
Tensor Tensor

A tensor of time steps.
Tensor Tensor

A tensor of time steps.
Source code in bionemo/moco/schedules/inference_time_schedules.py 
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
def generate_schedule(
    self,
    nsteps: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
) -> Tensor:
    """Generate the power time schedule as a tensor.

    Args:
        nsteps (Optional[int]): Number of time steps. If None uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").


    Returns:
        Tensor: A tensor of time steps.
        Tensor: A tensor of time steps.
    """
    if device is None:
        device = self.device
    if nsteps is None:
        nsteps = self.nsteps
    nsteps -= self.padding
    dilation = self.dilation + 1
    if dilation > 1:
        if nsteps % dilation != 0:
            raise ValueError(f"nsteps ({nsteps}) is not divisible by dilation + 1 ({dilation})")
        nsteps = int(nsteps / dilation)
    if nsteps is None:
        raise ValueError("nsteps cannot be None")
    if not self.inclusive_end:
        schedule = torch.linspace(self.min_t, 1, nsteps + 1).to(device=device) ** self.exponent
        schedule = schedule[:-1]
    else:
        schedule = torch.linspace(self.min_t, 1, nsteps).to(device=device) ** self.exponent
    if dilation > 1:
        schedule = schedule.repeat_interleave(dilation)
    if self.padding > 0:
        schedule = torch.cat((schedule, torch.ones(self.padding, device=device)))
    if self.direction == TimeDirection.DIFFUSION:
        schedule = 1 - schedule
    return schedule