109 lines
3.0 KiB
Python
109 lines
3.0 KiB
Python
from typing import Any, Dict, Tuple, Callable, Optional
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class PtnChainedModelDef(nn.Module):
|
|
"""
|
|
A chained model that sequentially applies `model_a` and `model_b`,
|
|
optionally followed by a closure function.
|
|
|
|
### Fields:
|
|
- `model_a`: First model in the chain.
|
|
- `model_b`: Second model in the chain.
|
|
- `closure`: Optional differentiable function to be called at the end of forward.
|
|
|
|
Example:
|
|
>>> model_a = nn.Linear(10, 20)
|
|
>>> model_b = nn.Linear(20, 5)
|
|
>>> closure = torch.nn.ReLU()
|
|
>>> model = PtnChainedModelDef(model_a, model_b, closure)
|
|
>>> x = torch.randn(1, 10)
|
|
>>> output = model(x)
|
|
|
|
pragma: skip_doc
|
|
"""
|
|
|
|
def __init__(self,
|
|
model_a: nn.Module,
|
|
model_b: nn.Module,
|
|
closure: Optional[Callable] = None) -> None:
|
|
super().__init__()
|
|
self.model_a = model_a
|
|
self.model_b = model_b
|
|
self.closure = closure
|
|
|
|
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Forward pass through the chained models.
|
|
|
|
Args:
|
|
inputs (torch.Tensor): Input tensor.
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after passing through `model_a`, `model_b`, and optionally `closure`.
|
|
"""
|
|
x = self.model_a(inputs)
|
|
x = self.model_b(x)
|
|
if self.closure:
|
|
x = self.closure(x)
|
|
return x
|
|
|
|
|
|
class PtnChainedModel:
|
|
"""
|
|
Constructs a chained PyTorch model.
|
|
|
|
### Fields:
|
|
- model_a: First model in the chain.
|
|
- model_b: Second model in the chain.
|
|
- closure: Optional differentiable function to be called at the end of forward.
|
|
|
|
category: PyTorch wrapper - Model
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the expected input types.
|
|
|
|
Returns:
|
|
Dict[str, Any]: A dictionary specifying required and optional input types.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"model_a": ("PTMODEL", {}),
|
|
"model_b": ("PTMODEL", {})
|
|
},
|
|
"optional": {
|
|
"closure": ("PTCALLABLE", {})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Training"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kw):
|
|
return float("NaN")
|
|
|
|
def f(self,
|
|
model_a: nn.Module,
|
|
model_b: nn.Module,
|
|
closure: Optional[Callable] = None) -> Tuple[nn.Module]:
|
|
"""
|
|
Constructs a chained model with optional closure.
|
|
|
|
Args:
|
|
model_a (nn.Module): First model in the chain.
|
|
model_b (nn.Module): Second model in the chain.
|
|
closure (Callable, optional): A function applied after `model_b`.
|
|
|
|
Returns:
|
|
Tuple[nn.Module]: A tuple containing the instantiated model.
|
|
"""
|
|
with torch.inference_mode(False):
|
|
model = PtnChainedModelDef(model_a, model_b, closure)
|
|
return (model,)
|