workflow.torch package

class workflow.torch.ModuleCompose(*modules_and_functions)[source]

ModuleCompose behaves like an extended torch.nn.Sequential that also allows:

  • vanilla functions

  • expands tuples to the next function’s arguments

  • specify a module and a wrapping function as seen in the example below

from torch import nn
from workflow.torch import ModuleCompose

ModuleCompose(
    nn.Conv2d(3, 32),
    nn.Conv2d(32, 16),
    (nn.Conv2d(16, 16), lambda conv, x: conv(x) + x),
    lambda x: (x, x * 2),
    lambda x, x2: x + x2,
)
forward(*x)[source]
debug(x)[source]
training: bool
workflow.torch.module_device(module)[source]
workflow.torch.set_seeds(seed)[source]
workflow.torch.module_train(module, training=True)

Can be used as a decorator or context manager

workflow.torch.module_eval(module, *, training=False)

Can be used as a decorator or context manager

workflow.torch.to_device(x, device)[source]
workflow.torch.get_model_summary(model)[source]