その他¶

その他のクラスと関数を提供します。

BaseModule¶

PyTorch LightningのLightningModuleを拡張した基底クラスです。

BaseModule ¶

Bases: LightningModule

Base module for PyTorch Lightning.

Functions¶

freeze_biases ¶

freeze_biases()

Freeze all bias parameters.

Source code in src/ml_networks/torch/base.py

def freeze_biases(self) -> None:
    """Freeze all bias parameters."""
    for name, param in self.named_parameters():
        if "bias" in name:
            param.requires_grad = False

freeze_weights ¶

freeze_weights()

Freeze all weight parameters.

Source code in src/ml_networks/torch/base.py

def freeze_weights(self) -> None:
    """Freeze all weight parameters."""
    for name, param in self.named_parameters():
        if "weight" in name:
            param.requires_grad = False

unfreeze_biases ¶

unfreeze_biases()

Unfreeze all bias parameters.

Source code in src/ml_networks/torch/base.py

def unfreeze_biases(self) -> None:
    """Unfreeze all bias parameters."""
    for name, param in self.named_parameters():
        if "bias" in name:
            param.requires_grad = True

unfreeze_weights ¶

unfreeze_weights()

Unfreeze all weight parameters.

Source code in src/ml_networks/torch/base.py

def unfreeze_weights(self) -> None:
    """Unfreeze all weight parameters."""
    for name, param in self.named_parameters():
        if "weight" in name:
            param.requires_grad = True

HyperNet¶

あるネットワーク（HyperNet）が別のネットワーク（TargetNet）の重みを動的に生成するメタ学習モジュールです。

HyperNet ¶

HyperNet(input_dim, output_params, fc_cfg=None, encoding=None)

Bases: LightningModule, HyperNetMixin

A hypernetwork that generates weights for a target network. Shape is dict[str, Tuple[int, ...]].

Args: input_dim: The dimension of the input. output_shapes: The shapes of the primary network weights being predicted. mlp_cfg: Configuration for the MLP backbone. encoding: The input encoding mode. Defaults to None.

Examples:

>>> from ml_networks.config import MLPConfig, LinearConfig
>>> input_dim = 10
>>> cond_dim = 128
>>> target_net = nn.Linear(10, 5)
>>> output_params = target_net.state_dict()
>>> mlp_cfg = MLPConfig(
...     hidden_dim=64,
...     n_layers=2,
...     output_activation="ReLU",
...     linear_cfg=LinearConfig(
...         activation="ReLU",
...     )
... )
>>> net = HyperNet(cond_dim, output_params, mlp_cfg)
>>> condition = torch.randn(2, cond_dim)
>>> x = torch.randn(2, input_dim)
>>> param = net(condition)
>>> outputs = torch.func.functional_call(target_net, param, x)
>>> outputs.shape
torch.Size([2, 5])

Source code in src/ml_networks/torch/hypernetworks.py

def __init__(
    self,
    input_dim: int,
    output_params: dict[str, Shape],
    fc_cfg: MLPConfig | None = None,
    encoding: InputMode = None,
) -> None:
    super().__init__()

    self.input_dim = input_dim
    self.output_shapes = output_params
    self.encoding = encoding

    # Cache this property to avoid recomputation
    self._output_offsets = self.output_offsets()

    self.backbone: nn.Module
    if fc_cfg is not None:
        self.backbone = MLPLayer(
            self.input_dim,
            self.flat_output_size(),
            fc_cfg,
        )
    else:
        self.backbone = nn.Linear(
            self.input_dim,
            self.flat_output_size(),
        )

Attributes¶

backbone `instance-attribute` ¶

backbone

encoding `instance-attribute` ¶

encoding = encoding

input_dim `instance-attribute` ¶

input_dim = input_dim

output_shapes `instance-attribute` ¶

output_shapes = output_params

Functions¶

forward ¶

forward(inputs)

Perform a forward pass of the neural network.

Args: inputs: The input tensors.

Returns:

Type	Description
`A dictionary of output tensors.`

Examples:

>>> from ml_networks.config import MLPConfig, LinearConfig
>>> input_dim = 10
>>> output_shapes = {"weight": (5, 10), "bias": (5,)}
>>> mlp_cfg = MLPConfig(
...     hidden_dim=64,
...     n_layers=2,
...     output_activation="ReLU",
...     linear_cfg=LinearConfig(
...         activation="ReLU",
...         norm="none",
...         dropout=0.0,
...         bias=True
...     )
... )
>>> net = HyperNet(input_dim, output_shapes, mlp_cfg)
>>> x = torch.randn(2, input_dim)
>>> outputs = net(x)
>>> outputs["weight"].shape
torch.Size([2, 5, 10])
>>> outputs["bias"].shape
torch.Size([2, 5])

Source code in src/ml_networks/torch/hypernetworks.py

def forward(self, inputs: torch.Tensor) -> dict[str, torch.Tensor]:
    """Perform a forward pass of the neural network.

    Args:
        inputs: The input tensors.

    Returns
    -------
        A dictionary of output tensors.

    Examples
    --------
    >>> from ml_networks.config import MLPConfig, LinearConfig
    >>> input_dim = 10
    >>> output_shapes = {"weight": (5, 10), "bias": (5,)}
    >>> mlp_cfg = MLPConfig(
    ...     hidden_dim=64,
    ...     n_layers=2,
    ...     output_activation="ReLU",
    ...     linear_cfg=LinearConfig(
    ...         activation="ReLU",
    ...         norm="none",
    ...         dropout=0.0,
    ...         bias=True
    ...     )
    ... )
    >>> net = HyperNet(input_dim, output_shapes, mlp_cfg)
    >>> x = torch.randn(2, input_dim)
    >>> outputs = net(x)
    >>> outputs["weight"].shape
    torch.Size([2, 5, 10])
    >>> outputs["bias"].shape
    torch.Size([2, 5])
    """
    if self.encoding is not None:
        inputs = encode_input(inputs, self.encoding)

    flat_output = self.backbone(inputs)
    return self.unflatten_output(flat_output)

ContrastiveLearningLoss¶

対照学習（Contrastive Learning）のための損失関数モジュールです。

ContrastiveLearningLoss ¶

ContrastiveLearningLoss(dim_input1, dim_input2, cfg)

Bases: LightningModule

Contrastive learning module.

Parameters:

Name	Type	Description	Default
`cfg`	`ContrastiveLearningConfig`	Configuration for contrastive learning.	required

Examples:

>>> from ml_networks.config import ContrastiveLearningConfig, MLPConfig, LinearConfig
>>> cfg = ContrastiveLearningConfig(
...     dim_feature=128,
...     dim_input1=256,
...     dim_input2=256,
...     eval_func=MLPConfig(
...         hidden_dim=256,
...         n_layers=2,
...         output_activation="ReLU",
...         linear_cfg=LinearConfig(
...             activation="ReLU",
...             norm="layer",
...             norm_cfg={"eps": 1e-05, "elementwise_affine": True, "bias": True},
...             dropout=0.1,
...             norm_first=False,
...             bias=True
...         )
...     ),
...     cross_entropy_like=False
... )
>>> model = ContrastiveLearningLoss(cfg)
>>> x1 = torch.randn(2, 256)
>>> x2 = torch.randn(2, 256)
>>> output = model.calc_nce(x1, x2)
>>> output["nce"].shape
torch.Size([])
>>> output, embeddings = model.calc_nce(x1, x2, return_emb=True)
>>> embeddings[0].shape, embeddings[1].shape
(torch.Size([2, 128]), torch.Size([2, 128]))

Source code in src/ml_networks/torch/contrastive.py

def __init__(
    self,
    dim_input1: int,
    dim_input2: int,
    cfg: ContrastiveLearningConfig,
) -> None:
    super().__init__()
    self.cfg = cfg
    self.dim_feature = cfg.dim_feature
    self.dim_input1 = dim_input1
    self.dim_input2 = dim_input2
    self.is_ce_like = cfg.cross_entropy_like

    self.eval_func = MLPLayer(dim_input1, cfg.dim_feature, cfg.eval_func)
    if self.dim_input1 != self.dim_input2:
        self.eval_func2 = MLPLayer(dim_input2, cfg.dim_feature, cfg.eval_func)
    else:
        self.eval_func2 = self.eval_func

Attributes¶

cfg `instance-attribute` ¶

cfg = cfg

dim_feature `instance-attribute` ¶

dim_feature = dim_feature

dim_input1 `instance-attribute` ¶

dim_input1 = dim_input1

dim_input2 `instance-attribute` ¶

dim_input2 = dim_input2

eval_func `instance-attribute` ¶

eval_func = MLPLayer(dim_input1, dim_feature, eval_func)

eval_func2 `instance-attribute` ¶

eval_func2 = MLPLayer(dim_input2, dim_feature, eval_func)

is_ce_like `instance-attribute` ¶

is_ce_like = cross_entropy_like

Functions¶

calc_nce ¶

calc_nce(feature1, feature2, return_emb=False)

Calculate the Noise Contrastive Estimation (NCE) loss.

Parameters:

Name	Type	Description	Default
`feature1`	`Tensor`	First input tensor of shape (*, dim_input1)	required
`feature2`	`Tensor`	Second input tensor of shape (*, dim_input2)	required
`return_emb`	`bool`	Whether to return embeddings, by default False	`False`

Returns:

Type	Description
`Union[dict[str, Tensor], Tuple[dict[str, Tensor], Tuple[Tensor, Tensor]]]`	If return_emb is False, returns loss dictionary. If return_emb is True, returns (loss dictionary, (embeddings1, embeddings2))

Source code in src/ml_networks/torch/contrastive.py

def calc_nce(
    self,
    feature1: torch.Tensor,
    feature2: torch.Tensor,
    return_emb: bool = False,
) -> dict[str, torch.Tensor] | tuple[dict[str, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
    """
    Calculate the Noise Contrastive Estimation (NCE) loss.

    Parameters
    ----------
    feature1 : torch.Tensor
        First input tensor of shape (*, dim_input1)
    feature2 : torch.Tensor
        Second input tensor of shape (*, dim_input2)
    return_emb : bool, optional
        Whether to return embeddings, by default False

    Returns
    -------
    Union[dict[str, torch.Tensor], Tuple[dict[str, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]]
        If return_emb is False, returns loss dictionary.
        If return_emb is True, returns (loss dictionary, (embeddings1, embeddings2))
    """
    loss_dict: dict[str, torch.Tensor] = {}
    batch_shape = feature1.shape[:-1]
    emb_1 = self.eval_func(feature1.reshape(-1, self.dim_input1))
    emb_2 = self.eval_func2(feature2.reshape(-1, self.dim_input2))

    if self.is_ce_like:
        labels = torch.arange(len(emb_1), device=emb_1.device)
        sim_matrix = torch.mm(emb_1, emb_2.T)
        nce_loss = F.cross_entropy(sim_matrix, labels, reduction="none") - np.log(len(sim_matrix))
        loss_dict["nce"] = nce_loss
    else:
        positive = torch.sum(emb_1 * emb_2, dim=-1)
        loss_dict["positive"] = positive.detach().clone().mean()

        sim_matrix = torch.mm(emb_1, emb_2.T)
        negative = torch.logsumexp(sim_matrix, dim=-1) - np.log(len(sim_matrix))
        loss_dict["negative"] = negative.detach().clone().mean()

        nce_loss = -positive + negative
        loss_dict["nce"] = nce_loss.reshape(batch_shape)

    if return_emb:
        return loss_dict, (emb_1, emb_2)
    return loss_dict

calc_sigmoid ¶

calc_sigmoid(feature1, feature2, return_emb=False, temperature=0.1, bias=0.0)

Calculate the Sigmoid loss for contrastive learning.

Parameters:

Name	Type	Description	Default
`feature1`	`Tensor`	First input tensor of shape (*, dim_input1)	required
`feature2`	`Tensor`	Second input tensor of shape (*, dim_input2)	required
`return_emb`	`bool`	Whether to return embeddings, by default False	`False`

Returns:

Type	Description
`Union[dict[str, Tensor], Tuple[dict[str, Tensor], Tuple[Tensor, Tensor]]]`	If return_emb is False, returns loss dictionary. If return_emb is True, returns (loss dictionary, (embeddings1, embeddings2))

Source code in src/ml_networks/torch/contrastive.py

def calc_sigmoid(
    self,
    feature1: torch.Tensor,
    feature2: torch.Tensor,
    return_emb: bool = False,
    temperature: float | torch.Tensor = 0.1,
    bias: float | torch.Tensor = 0.0,
) -> dict[str, torch.Tensor] | tuple[dict[str, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
    """
    Calculate the Sigmoid loss for contrastive learning.

    Parameters
    ----------
    feature1 : torch.Tensor
        First input tensor of shape (*, dim_input1)
    feature2 : torch.Tensor
        Second input tensor of shape (*, dim_input2)
    return_emb : bool, optional
        Whether to return embeddings, by default False

    Returns
    -------
    Union[dict[str, torch.Tensor], Tuple[dict[str, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]]
        If return_emb is False, returns loss dictionary.
        If return_emb is True, returns (loss dictionary, (embeddings1, embeddings2))
    """
    loss_dict: dict[str, torch.Tensor] = {}
    batch_shape = feature1.shape[:-1]
    emb_1 = self.eval_func(feature1.reshape(-1, self.dim_input1))
    emb_2 = self.eval_func2(feature2.reshape(-1, self.dim_input2))

    logits = torch.matmul(emb_1, emb_2.T) * temperature + bias
    labels = torch.eye(len(logits), device=logits.device) * 2 - 1
    loss = -F.logsigmoid(logits * labels).sum(-1)
    loss_dict["sigmoid"] = loss.reshape(batch_shape)
    if return_emb:
        return loss_dict, (emb_1, emb_2)
    return loss_dict

calc_timeseries_nce ¶

calc_timeseries_nce(feature1, feature2, positive_range_self=0, positive_range_tgt=0, return_emb=False)

Calculate the Noise Contrastive Estimation (NCE) loss for time series data.

Parameters:

Name	Type	Description	Default
`feature1`	`Tensor`	First input tensor of shape (*batch, length, dim_input1)	required
`feature2`	`Tensor`	Second input tensor of shape (*batch, length, dim_input2)	required
`positive_range_self`	`int`	Range for self-positive samples, by default 0	`0`
`positive_range_tgt`	`int`	Range for target-positive samples, by default 0	`0`
`return_emb`	`bool`	Whether to return embeddings, by default False	`False`

Returns:

Type	Description
`Union[dict[str, Tensor], Tuple[dict[str, Tensor], Tuple[Tensor, Tensor]]]`	If return_emb is False, returns loss dictionary. If return_emb is True, returns (loss dictionary, (embeddings1, embeddings2))

Source code in src/ml_networks/torch/contrastive.py

def calc_timeseries_nce(
    self,
    feature1: torch.Tensor,
    feature2: torch.Tensor,
    positive_range_self: int = 0,
    positive_range_tgt: int = 0,
    return_emb: bool = False,
) -> dict[str, torch.Tensor] | tuple[dict[str, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
    """
    Calculate the Noise Contrastive Estimation (NCE) loss for time series data.

    Parameters
    ----------
    feature1 : torch.Tensor
        First input tensor of shape (*batch, length, dim_input1)
    feature2 : torch.Tensor
        Second input tensor of shape (*batch, length, dim_input2)
    positive_range_self : int, optional
        Range for self-positive samples, by default 0
    positive_range_tgt : int, optional
        Range for target-positive samples, by default 0
    return_emb : bool, optional
        Whether to return embeddings, by default False

    Returns
    -------
    Union[dict[str, torch.Tensor], Tuple[dict[str, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]]
        If return_emb is False, returns loss dictionary.
        If return_emb is True, returns (loss dictionary, (embeddings1, embeddings2))
    """
    if not positive_range_self and not positive_range_tgt:
        return self.calc_nce(feature1, feature2, return_emb)

    # Reshape inputs
    feature1 = feature1.reshape(-1, feature1.shape[-2], self.dim_input1)
    feature2 = feature2.reshape(-1, feature2.shape[-2], self.dim_input2)
    batch, length, _ = feature1.shape

    # Calculate embeddings
    emb_1 = self.eval_func(feature1)
    emb_2 = self.eval_func2(feature2)

    # Initialize loss dictionary
    loss_dict: dict[str, torch.Tensor] = {}

    # Calculate positive pairs
    positive = torch.sum(emb_1.flatten(0, 1) * emb_2.flatten(0, 1), dim=-1)  # (batch*length)
    loss_dict["positive"] = positive.detach().clone().mean()

    # Calculate self-positive pairs if needed
    if positive_range_self > 0:
        self_positive_1, self_positive_2 = self._calculate_self_positive_pairs(
            emb_1,
            emb_2,
            batch,
            length,
            positive_range_self,
        )
        positive += self_positive_1.flatten(0, 1) + self_positive_2.flatten(0, 1)
        loss_dict["self_positive_1"] = self_positive_1.detach().clone().mean()
        loss_dict["self_positive_2"] = self_positive_2.detach().clone().mean()

    # Calculate target-positive pairs if needed
    if positive_range_tgt > 0:
        tgt_positive = self._calculate_target_positive_pairs(
            emb_1,
            emb_2,
            batch,
            length,
            positive_range_tgt,
        )
        positive += tgt_positive.flatten(0, 1)
        loss_dict["tgt_positive"] = tgt_positive.detach().clone().mean()

    # Calculate negative pairs
    sim_matrix = torch.mm(emb_1.flatten(0, 1), emb_2.flatten(0, 1).T)
    negative = torch.logsumexp(sim_matrix, dim=-1) - np.log(len(sim_matrix))
    loss_dict["negative"] = negative.detach().clone().mean()

    # Calculate final loss
    nce_loss = -positive + negative
    nce_loss = nce_loss.mean()
    loss_dict["nce"] = nce_loss

    if return_emb:
        return loss_dict, (emb_1, emb_2)
    return loss_dict

ProgressBarCallback¶

PyTorch LightningのRichプログレスバーコールバックです。

ProgressBarCallback ¶

ProgressBarCallback()

Bases: RichProgressBar

Make the progress bar richer.

References

https://qiita.com/akihironitta/items/edfd6b29dfb67b17fb00

Rich progress bar with custom theme.

Source code in src/ml_networks/callbacks.py

def __init__(self) -> None:
    """Rich progress bar with custom theme."""
    theme = RichProgressBarTheme(
        description="green_yellow",
        progress_bar="green1",
        progress_bar_finished="green1",
        batch_progress="green_yellow",
        time="grey82",
        processing_speed="grey82",
        metrics="grey82",
    )
    super().__init__(theme=theme)

その他¶

BaseModule¶

BaseModule ¶

Functions¶

freeze_biases ¶

freeze_weights ¶

unfreeze_biases ¶

unfreeze_weights ¶

HyperNet¶

HyperNet ¶

Attributes¶

backbone instance-attribute ¶

encoding instance-attribute ¶

input_dim instance-attribute ¶

output_shapes instance-attribute ¶

Functions¶

forward ¶

ContrastiveLearningLoss¶

ContrastiveLearningLoss ¶

Attributes¶

cfg instance-attribute ¶

dim_feature instance-attribute ¶

dim_input1 instance-attribute ¶

dim_input2 instance-attribute ¶

eval_func instance-attribute ¶

eval_func2 instance-attribute ¶

is_ce_like instance-attribute ¶

Functions¶

calc_nce ¶

calc_sigmoid ¶

calc_timeseries_nce ¶

ProgressBarCallback¶

ProgressBarCallback ¶

Functions¶

backbone `instance-attribute` ¶

encoding `instance-attribute` ¶

input_dim `instance-attribute` ¶

output_shapes `instance-attribute` ¶

cfg `instance-attribute` ¶

dim_feature `instance-attribute` ¶

dim_input1 `instance-attribute` ¶

dim_input2 `instance-attribute` ¶

eval_func `instance-attribute` ¶

eval_func2 `instance-attribute` ¶

is_ce_like `instance-attribute` ¶