workflow.torch package

class workflow.torch.ModuleCompose(*modules_and_functions)

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)

debug(x)

training: bool

workflow.torch.module_device(module)

workflow.torch.set_seeds(seed)

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)

workflow.torch.get_model_summary(model)