Spaces:
Runtime error
Runtime error
File size: 1,334 Bytes
cc0dd3c |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, List
import torch
from mmpretrain.registry import MODELS
from mmpretrain.structures import DataSample
from .base import BaseSelfSupervisor
@MODELS.register_module()
class SimSiam(BaseSelfSupervisor):
"""SimSiam.
Implementation of `Exploring Simple Siamese Representation Learning
<https://arxiv.org/abs/2011.10566>`_. The operation of fixing learning rate
of predictor is in `engine/hooks/simsiam_hook.py`.
"""
def loss(self, inputs: List[torch.Tensor], data_samples: List[DataSample],
**kwargs) -> Dict[str, torch.Tensor]:
"""The forward function in training.
Args:
inputs (List[torch.Tensor]): The input images.
data_samples (List[DataSample]): All elements required
during the forward function.
Returns:
Dict[str, torch.Tensor]: A dictionary of loss components.
"""
assert isinstance(inputs, list)
img_v1 = inputs[0]
img_v2 = inputs[1]
z1 = self.neck(self.backbone(img_v1))[0] # NxC
z2 = self.neck(self.backbone(img_v2))[0] # NxC
loss_1 = self.head.loss(z1, z2)
loss_2 = self.head.loss(z2, z1)
losses = dict(loss=0.5 * (loss_1 + loss_2))
return losses
|