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
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
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.
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
- extra_repr
- compile