50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
from typing import Any, Dict, Tuple
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
class PtComputeLoss:
|
|
"""
|
|
Computes loss for the input with the target using the specified loss function.
|
|
|
|
The user does not normally need to use this node. Instead the user should use training nodes that accept a loss object.
|
|
|
|
category: PyTorch wrapper - Training
|
|
"""
|
|
|
|
@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": {
|
|
"input_tens": ("TENSOR", {}),
|
|
"target_tens": ("TENSOR", {}),
|
|
"loss": ("PTLOSS", {})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: Tuple[str, ...] = ("TENSOR",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Training"
|
|
|
|
def f(self,
|
|
input_tens: torch.Tensor,
|
|
target_tens: torch.Tensor,
|
|
loss: nn.Module) -> Tuple[torch.Tensor]:
|
|
"""
|
|
Computes loss for the input with the target using the specified loss function.
|
|
|
|
Args:
|
|
input_tens (torch.Tensor): Tensor containing the model output or y_hat.
|
|
target_tens (torch.Tensor): Tensor containing the ground truth or the label.
|
|
loss (nn.Module): Loss object instance.
|
|
|
|
Returns:
|
|
Tuple[nn.Module]: A tuple containing the computed loss.
|
|
"""
|
|
return (loss(input_tens, target_tens),)
|