File size: 2,865 Bytes
26225c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | from torch import nn
__all__ = ['MultiLoss']
class MultiLoss(nn.Module):
"""Wrapper to compute the weighted sum of multiple criteria
:param criteria: List(callable)
List of criteria
:param lambdas: List(str)
"""
def __init__(self, criteria, lambdas):
super().__init__()
assert len(criteria) == len(lambdas)
self.criteria = nn.ModuleList(criteria)
self.lambdas = lambdas
def __len__(self):
return len(self.criteria)
def to(self, *args, **kwargs):
for i in range(len(self)):
self.criteria[i] = self.criteria[i].to(*args, **kwargs)
self.lambdas[i] = self.lambdas[i].to(*args, **kwargs)
def extra_repr(self) -> str:
return f'lambdas={self.lambdas}'
def forward(self, a, b, **kwargs):
loss = 0
for lamb, criterion, a_, b_ in zip(self.lambdas, self.criteria, a, b):
loss = loss + lamb * criterion(a_, b_, **kwargs)
return loss
@property
def weight(self):
"""MultiLoss supports `weight` if all its criteria support it.
"""
return self.criteria[0].weight
@weight.setter
def weight(self, weight):
"""MultiLoss supports `weight` if all its criteria support it.
"""
for i in range(len(self)):
self.criteria[i].weight = weight
def state_dict(self, *args, destination=None, prefix='', keep_vars=False):
"""Normal `state_dict` behavior, except for the shared criterion
weights, which are not saved under `prefix.criteria.i.weight`
but under `prefix.weight`.
"""
destination = super().state_dict(
*args, destination=destination, prefix=prefix, keep_vars=keep_vars)
# Remove the 'weight' from the criteria
for i in range(len(self)):
destination.pop(f"{prefix}criteria.{i}.weight")
# Only save the global shared weight
destination[f"{prefix}weight"] = self.weight
return destination
def load_state_dict(self, state_dict, strict=True):
"""Normal `load_state_dict` behavior, except for the shared
criterion weights, which are not saved under `criteria.i.weight`
but under `prefix.weight`.
"""
# Get the weight from the state_dict
old_format = state_dict.get('criteria.0.weight')
new_format = state_dict.get('weight')
weight = new_format if new_format is not None else old_format
for k in [f"criteria.{i}.weight" for i in range(len(self))]:
if k in state_dict.keys():
state_dict.pop(k)
# Normal load_state_dict, ignoring self.criteria.0.weight and
# self.weight
out = super().load_state_dict(state_dict, strict=strict)
# Set the weight
self.weight = weight
return out
|