torch_em.loss.combined_loss 

 1import torch.nn as nn
 2
 3
 4class CombinedLoss(nn.Module):
 5    def __init__(self, *losses, loss_weights=None):
 6        super().__init__()
 7        self.losses = nn.ModuleList(losses)
 8        n_losses = len(self.losses)
 9        if loss_weights is None:
10            try:
11                self.loss_weights = [1.0 / n_losses] * n_losses
12            except ZeroDivisionError:
13                self.loss_weights = None
14        else:
15            assert len(loss_weights) == n_losses
16            self.loss_weights = loss_weights
17
18    def forward(self, x, y):
19        assert self.loss_weights is not None
20        loss_value = sum([loss(x, y) * weight for loss, weight in zip(self.losses, self.loss_weights)])
21        return loss_value
class CombinedLoss(torch.nn.modules.module.Module): 
 5class CombinedLoss(nn.Module):
 6    def __init__(self, *losses, loss_weights=None):
 7        super().__init__()
 8        self.losses = nn.ModuleList(losses)
 9        n_losses = len(self.losses)
10        if loss_weights is None:
11            try:
12                self.loss_weights = [1.0 / n_losses] * n_losses
13            except ZeroDivisionError:
14                self.loss_weights = None
15        else:
16            assert len(loss_weights) == n_losses
17            self.loss_weights = loss_weights
18
19    def forward(self, x, y):
20        assert self.loss_weights is not None
21        loss_value = sum([loss(x, y) * weight for loss, weight in zip(self.losses, self.loss_weights)])
22        return loss_value

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

CombinedLoss(*losses, loss_weights=None) 
 6    def __init__(self, *losses, loss_weights=None):
 7        super().__init__()
 8        self.losses = nn.ModuleList(losses)
 9        n_losses = len(self.losses)
10        if loss_weights is None:
11            try:
12                self.loss_weights = [1.0 / n_losses] * n_losses
13            except ZeroDivisionError:
14                self.loss_weights = None
15        else:
16            assert len(loss_weights) == n_losses
17            self.loss_weights = loss_weights

Initializes internal Module state, shared by both nn.Module and ScriptModule.

losses
def forward(self, x, y): 
19    def forward(self, x, y):
20        assert self.loss_weights is not None
21        loss_value = sum([loss(x, y) * weight for loss, weight in zip(self.losses, self.loss_weights)])
22        return loss_value

Defines the computation performed at every call.

Should be overridden by all subclasses.

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Inherited Members
torch.nn.modules.module.Module
dump_patches
training
call_super_init
register_buffer
register_parameter
add_module
register_module
get_submodule
get_parameter
get_buffer
get_extra_state
set_extra_state
apply
cuda
ipu
xpu
cpu
type
float
double
half
bfloat16
to_empty
to
register_full_backward_pre_hook
register_backward_hook
register_full_backward_hook
register_forward_pre_hook
register_forward_hook
register_state_dict_pre_hook
state_dict
register_load_state_dict_post_hook
load_state_dict
parameters
named_parameters
buffers
named_buffers
children
named_children
modules
named_modules
train
eval
requires_grad_
zero_grad
share_memory
extra_repr
compile