add projects code

projects/maskfeat_video/models/maskfeat.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List
+
+import torch
+import torch.nn.functional as F
+
+from mmpretrain.models import BaseSelfSupervisor
+from mmpretrain.registry import MODELS
+from mmpretrain.structures import DataSample
+
+
+@MODELS.register_module()
+class VideoMaskFeat(BaseSelfSupervisor):
+    """MaskFeat.
+
+    Implementation of `Masked Feature Prediction for Self-Supervised Visual
+    Pre-Training <https://arxiv.org/abs/2112.09133>`_.
+    """
+
+    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.
+        """
+        mask = torch.stack(
+            [data_sample.mask.value for data_sample in data_samples])
+        mask = mask.to(torch.bool)
+
+        video = inputs[0]
+        video = video.view((-1, ) + video.shape[2:])  # B, C, T, H, W
+        latent = self.backbone(video, mask)
+        B, L, C = latent[0].shape
+        pred = self.neck([latent[0].view(B * L, C)])
+        pred = pred[0].view(B, L, -1)
+
+        # generate hog target
+        video = video[:, :, ::self.backbone.patch_stride[0], :, :]
+        video = video.transpose(1, 2)  # B, T, C, H, W
+        self.target_generator.B = video.size(0)
+        self.target_generator.T = video.size(1)
+        video = video.flatten(0, 1)  # B*T, C, H, W
+        hog = self.target_generator(video)
+
+        mask = self._get_output_mask(mask)
+        loss = self.head(pred, hog, mask)
+        losses = dict(loss=loss)
+        return losses
+
+    def _get_output_mask(self, mask: torch.Tensor) -> torch.Tensor:
+        size = self.backbone.out_patch_resolution[-1][-1]
+        output_mask = F.interpolate(mask.float(), size=size)
+        return output_mask

projects/maskfeat_video/tools/dist_train.sh

+#!/usr/bin/env bash
+
+CONFIG=$1
+GPUS=$2
+NNODES=${NNODES:-1}
+NODE_RANK=${NODE_RANK:-0}
+PORT=${PORT:-29500}
+MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+
+PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
+python -m torch.distributed.launch \
+    --nnodes=$NNODES \
+    --node_rank=$NODE_RANK \
+    --master_addr=$MASTER_ADDR \
+    --nproc_per_node=$GPUS \
+    --master_port=$PORT \
+    $(dirname "$0")/train.py \
+    $CONFIG \
+    --launcher pytorch ${@:3}

projects/maskfeat_video/tools/slurm_train.sh

+#!/usr/bin/env bash
+
+set -x
+
+PARTITION=$1
+JOB_NAME=$2
+CONFIG=$3
+GPUS=${GPUS:-8}
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+CPUS_PER_TASK=${CPUS_PER_TASK:-5}
+SRUN_ARGS=${SRUN_ARGS:-""}
+PY_ARGS=${@:4}
+
+PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
+srun -p ${PARTITION} \
+    --job-name=${JOB_NAME} \
+    --gres=gpu:${GPUS_PER_NODE} \
+    --ntasks=${GPUS} \
+    --ntasks-per-node=${GPUS_PER_NODE} \
+    --cpus-per-task=${CPUS_PER_TASK} \
+    --kill-on-bad-exit=1 \
+    ${SRUN_ARGS} \
+    python -u tools/train.py ${CONFIG} --launcher="slurm" ${PY_ARGS}