Initial media depth project backup
This commit is contained in:
Codex
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .train import get_args_parser, main
|
||||
from .ssl_meta_arch import SSLMetaArch
|
||||
@@ -0,0 +1,403 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from functools import partial
|
||||
import logging
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from dinov2.loss import DINOLoss, iBOTPatchLoss, KoLeoLoss
|
||||
from dinov2.models import build_model_from_cfg
|
||||
from dinov2.layers import DINOHead
|
||||
from dinov2.utils.utils import has_batchnorms
|
||||
from dinov2.utils.param_groups import get_params_groups_with_decay, fuse_params_groups
|
||||
from dinov2.fsdp import get_fsdp_wrapper, ShardedGradScaler, get_fsdp_modules, reshard_fsdp_model
|
||||
|
||||
from dinov2.models.vision_transformer import BlockChunk
|
||||
|
||||
try:
|
||||
from xformers.ops import fmha
|
||||
|
||||
XFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
XFORMERS_AVAILABLE = False
|
||||
assert XFORMERS_AVAILABLE, "xFormers is required for DINOv2 training"
|
||||
|
||||
|
||||
logger = logging.getLogger("dinov2")
|
||||
|
||||
|
||||
class SSLMetaArch(nn.Module):
|
||||
def __init__(self, cfg):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
self.fp16_scaler = ShardedGradScaler() if cfg.compute_precision.grad_scaler else None
|
||||
|
||||
student_model_dict = dict()
|
||||
teacher_model_dict = dict()
|
||||
|
||||
student_backbone, teacher_backbone, embed_dim = build_model_from_cfg(cfg)
|
||||
student_model_dict["backbone"] = student_backbone
|
||||
teacher_model_dict["backbone"] = teacher_backbone
|
||||
logger.info(f"OPTIONS -- architecture : embed_dim: {embed_dim}")
|
||||
|
||||
if cfg.student.pretrained_weights:
|
||||
chkpt = torch.load(cfg.student.pretrained_weights)
|
||||
logger.info(f"OPTIONS -- pretrained weights: loading from {cfg.student.pretrained_weights}")
|
||||
student_backbone.load_state_dict(chkpt["model"], strict=False)
|
||||
|
||||
self.embed_dim = embed_dim
|
||||
self.dino_out_dim = cfg.dino.head_n_prototypes
|
||||
|
||||
self.do_dino = cfg.dino.loss_weight > 0
|
||||
self.do_koleo = cfg.dino.koleo_loss_weight > 0
|
||||
self.do_ibot = cfg.ibot.loss_weight > 0
|
||||
self.ibot_separate_head = cfg.ibot.separate_head
|
||||
|
||||
logger.info("OPTIONS -- DINO")
|
||||
if self.do_dino:
|
||||
logger.info(f"OPTIONS -- DINO -- loss_weight: {cfg.dino.loss_weight}")
|
||||
logger.info(f"OPTIONS -- DINO -- head_n_prototypes: {cfg.dino.head_n_prototypes}")
|
||||
logger.info(f"OPTIONS -- DINO -- head_bottleneck_dim: {cfg.dino.head_bottleneck_dim}")
|
||||
logger.info(f"OPTIONS -- DINO -- head_hidden_dim: {cfg.dino.head_hidden_dim}")
|
||||
self.dino_loss_weight = cfg.dino.loss_weight
|
||||
dino_head = partial(
|
||||
DINOHead,
|
||||
in_dim=embed_dim,
|
||||
out_dim=cfg.dino.head_n_prototypes,
|
||||
hidden_dim=cfg.dino.head_hidden_dim,
|
||||
bottleneck_dim=cfg.dino.head_bottleneck_dim,
|
||||
nlayers=cfg.dino.head_nlayers,
|
||||
)
|
||||
self.dino_loss = DINOLoss(self.dino_out_dim)
|
||||
if self.do_koleo:
|
||||
logger.info("OPTIONS -- DINO -- applying KOLEO regularization")
|
||||
self.koleo_loss = KoLeoLoss()
|
||||
|
||||
else:
|
||||
logger.info("OPTIONS -- DINO -- not using DINO")
|
||||
|
||||
if self.do_dino or self.do_ibot:
|
||||
student_model_dict["dino_head"] = dino_head()
|
||||
teacher_model_dict["dino_head"] = dino_head()
|
||||
|
||||
logger.info("OPTIONS -- IBOT")
|
||||
logger.info(f"OPTIONS -- IBOT -- loss_weight: {cfg.ibot.loss_weight}")
|
||||
logger.info(f"OPTIONS -- IBOT masking -- ibot_mask_ratio_tuple: {cfg.ibot.mask_ratio_min_max}")
|
||||
logger.info(f"OPTIONS -- IBOT masking -- ibot_mask_sample_probability: {cfg.ibot.mask_sample_probability}")
|
||||
if self.do_ibot:
|
||||
self.ibot_loss_weight = cfg.ibot.loss_weight
|
||||
assert max(cfg.ibot.mask_ratio_min_max) > 0, "please provide a positive mask ratio tuple for ibot"
|
||||
assert cfg.ibot.mask_sample_probability > 0, "please provide a positive mask probability for ibot"
|
||||
self.ibot_out_dim = cfg.ibot.head_n_prototypes if self.ibot_separate_head else cfg.dino.head_n_prototypes
|
||||
self.ibot_patch_loss = iBOTPatchLoss(self.ibot_out_dim)
|
||||
if self.ibot_separate_head:
|
||||
logger.info(f"OPTIONS -- IBOT -- loss_weight: {cfg.ibot.loss_weight}")
|
||||
logger.info(f"OPTIONS -- IBOT -- head_n_prototypes: {cfg.ibot.head_n_prototypes}")
|
||||
logger.info(f"OPTIONS -- IBOT -- head_bottleneck_dim: {cfg.ibot.head_bottleneck_dim}")
|
||||
logger.info(f"OPTIONS -- IBOT -- head_hidden_dim: {cfg.ibot.head_hidden_dim}")
|
||||
ibot_head = partial(
|
||||
DINOHead,
|
||||
in_dim=embed_dim,
|
||||
out_dim=cfg.ibot.head_n_prototypes,
|
||||
hidden_dim=cfg.ibot.head_hidden_dim,
|
||||
bottleneck_dim=cfg.ibot.head_bottleneck_dim,
|
||||
nlayers=cfg.ibot.head_nlayers,
|
||||
)
|
||||
student_model_dict["ibot_head"] = ibot_head()
|
||||
teacher_model_dict["ibot_head"] = ibot_head()
|
||||
else:
|
||||
logger.info("OPTIONS -- IBOT -- head shared with DINO")
|
||||
|
||||
self.need_to_synchronize_fsdp_streams = True
|
||||
|
||||
self.student = nn.ModuleDict(student_model_dict)
|
||||
self.teacher = nn.ModuleDict(teacher_model_dict)
|
||||
|
||||
# there is no backpropagation through the teacher, so no need for gradients
|
||||
for p in self.teacher.parameters():
|
||||
p.requires_grad = False
|
||||
logger.info(f"Student and Teacher are built: they are both {cfg.student.arch} network.")
|
||||
|
||||
def forward(self, inputs):
|
||||
raise NotImplementedError
|
||||
|
||||
def backprop_loss(self, loss):
|
||||
if self.fp16_scaler is not None:
|
||||
self.fp16_scaler.scale(loss).backward()
|
||||
else:
|
||||
loss.backward()
|
||||
|
||||
def forward_backward(self, images, teacher_temp):
|
||||
n_global_crops = 2
|
||||
assert n_global_crops == 2
|
||||
n_local_crops = self.cfg.crops.local_crops_number
|
||||
|
||||
global_crops = images["collated_global_crops"].cuda(non_blocking=True)
|
||||
local_crops = images["collated_local_crops"].cuda(non_blocking=True)
|
||||
|
||||
masks = images["collated_masks"].cuda(non_blocking=True)
|
||||
mask_indices_list = images["mask_indices_list"].cuda(non_blocking=True)
|
||||
n_masked_patches_tensor = images["n_masked_patches"].cuda(non_blocking=True)
|
||||
n_masked_patches = mask_indices_list.shape[0]
|
||||
upperbound = images["upperbound"]
|
||||
masks_weight = images["masks_weight"].cuda(non_blocking=True)
|
||||
|
||||
n_local_crops_loss_terms = max(n_local_crops * n_global_crops, 1)
|
||||
n_global_crops_loss_terms = (n_global_crops - 1) * n_global_crops
|
||||
|
||||
do_dino = self.do_dino
|
||||
do_ibot = self.do_ibot
|
||||
|
||||
# loss scales
|
||||
ibot_loss_scale = 1.0 / n_global_crops
|
||||
|
||||
# teacher output
|
||||
@torch.no_grad()
|
||||
def get_teacher_output():
|
||||
x, n_global_crops_teacher = global_crops, n_global_crops
|
||||
teacher_backbone_output_dict = self.teacher.backbone(x, is_training=True)
|
||||
teacher_cls_tokens = teacher_backbone_output_dict["x_norm_clstoken"]
|
||||
teacher_cls_tokens = teacher_cls_tokens.chunk(n_global_crops_teacher)
|
||||
# watch out: these are chunked and cat'd in reverse so A is matched to B in the global crops dino loss
|
||||
teacher_cls_tokens = torch.cat((teacher_cls_tokens[1], teacher_cls_tokens[0]))
|
||||
ibot_teacher_patch_tokens = teacher_backbone_output_dict["x_norm_patchtokens"]
|
||||
_dim = ibot_teacher_patch_tokens.shape[-1]
|
||||
n_cls_tokens = teacher_cls_tokens.shape[0]
|
||||
|
||||
if do_ibot and not self.ibot_separate_head:
|
||||
buffer_tensor_teacher = ibot_teacher_patch_tokens.new_zeros(upperbound + n_cls_tokens, _dim)
|
||||
buffer_tensor_teacher[:n_cls_tokens].copy_(teacher_cls_tokens)
|
||||
torch.index_select(
|
||||
ibot_teacher_patch_tokens.flatten(0, 1),
|
||||
dim=0,
|
||||
index=mask_indices_list,
|
||||
out=buffer_tensor_teacher[n_cls_tokens : n_cls_tokens + n_masked_patches],
|
||||
)
|
||||
tokens_after_head = self.teacher.dino_head(buffer_tensor_teacher)
|
||||
teacher_cls_tokens_after_head = tokens_after_head[:n_cls_tokens]
|
||||
masked_teacher_patch_tokens_after_head = tokens_after_head[
|
||||
n_cls_tokens : n_cls_tokens + n_masked_patches
|
||||
]
|
||||
elif do_ibot and self.ibot_separate_head:
|
||||
buffer_tensor_teacher = ibot_teacher_patch_tokens.new_zeros(upperbound, _dim)
|
||||
torch.index_select(
|
||||
ibot_teacher_patch_tokens.flatten(0, 1),
|
||||
dim=0,
|
||||
index=mask_indices_list,
|
||||
out=buffer_tensor_teacher[:n_masked_patches],
|
||||
)
|
||||
teacher_cls_tokens_after_head = self.teacher.dino_head(teacher_cls_tokens)
|
||||
masked_teacher_patch_tokens_after_head = self.teacher.ibot_head(buffer_tensor_teacher)[
|
||||
:n_masked_patches
|
||||
]
|
||||
else:
|
||||
teacher_cls_tokens_after_head = self.teacher.dino_head(teacher_cls_tokens)
|
||||
masked_teacher_ibot_softmaxed_centered = None
|
||||
|
||||
if self.cfg.train.centering == "centering":
|
||||
teacher_dino_softmaxed_centered_list = self.dino_loss.softmax_center_teacher(
|
||||
teacher_cls_tokens_after_head, teacher_temp=teacher_temp
|
||||
).view(n_global_crops_teacher, -1, *teacher_cls_tokens_after_head.shape[1:])
|
||||
self.dino_loss.update_center(teacher_cls_tokens_after_head)
|
||||
if do_ibot:
|
||||
masked_teacher_patch_tokens_after_head = masked_teacher_patch_tokens_after_head.unsqueeze(0)
|
||||
masked_teacher_ibot_softmaxed_centered = self.ibot_patch_loss.softmax_center_teacher(
|
||||
masked_teacher_patch_tokens_after_head[:, :n_masked_patches], teacher_temp=teacher_temp
|
||||
)
|
||||
masked_teacher_ibot_softmaxed_centered = masked_teacher_ibot_softmaxed_centered.squeeze(0)
|
||||
self.ibot_patch_loss.update_center(masked_teacher_patch_tokens_after_head[:n_masked_patches])
|
||||
|
||||
elif self.cfg.train.centering == "sinkhorn_knopp":
|
||||
teacher_dino_softmaxed_centered_list = self.dino_loss.sinkhorn_knopp_teacher(
|
||||
teacher_cls_tokens_after_head, teacher_temp=teacher_temp
|
||||
).view(n_global_crops_teacher, -1, *teacher_cls_tokens_after_head.shape[1:])
|
||||
|
||||
if do_ibot:
|
||||
masked_teacher_ibot_softmaxed_centered = self.ibot_patch_loss.sinkhorn_knopp_teacher(
|
||||
masked_teacher_patch_tokens_after_head,
|
||||
teacher_temp=teacher_temp,
|
||||
n_masked_patches_tensor=n_masked_patches_tensor,
|
||||
)
|
||||
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return teacher_dino_softmaxed_centered_list, masked_teacher_ibot_softmaxed_centered
|
||||
|
||||
teacher_dino_softmaxed_centered_list, masked_teacher_ibot_softmaxed_centered = get_teacher_output()
|
||||
reshard_fsdp_model(self.teacher)
|
||||
|
||||
loss_dict = {}
|
||||
|
||||
loss_accumulator = 0 # for backprop
|
||||
student_global_backbone_output_dict, student_local_backbone_output_dict = self.student.backbone(
|
||||
[global_crops, local_crops], masks=[masks, None], is_training=True
|
||||
)
|
||||
|
||||
inputs_for_student_head_list = []
|
||||
|
||||
# 1a: local crops cls tokens
|
||||
student_local_cls_tokens = student_local_backbone_output_dict["x_norm_clstoken"]
|
||||
inputs_for_student_head_list.append(student_local_cls_tokens.unsqueeze(0))
|
||||
|
||||
# 1b: global crops cls tokens
|
||||
student_global_cls_tokens = student_global_backbone_output_dict["x_norm_clstoken"]
|
||||
inputs_for_student_head_list.append(student_global_cls_tokens.unsqueeze(0))
|
||||
|
||||
# 1c: global crops patch tokens
|
||||
if do_ibot:
|
||||
_dim = student_global_backbone_output_dict["x_norm_clstoken"].shape[-1]
|
||||
ibot_student_patch_tokens = student_global_backbone_output_dict["x_norm_patchtokens"]
|
||||
buffer_tensor_patch_tokens = ibot_student_patch_tokens.new_zeros(upperbound, _dim)
|
||||
buffer_tensor_patch_tokens[:n_masked_patches].copy_(
|
||||
torch.index_select(ibot_student_patch_tokens.flatten(0, 1), dim=0, index=mask_indices_list)
|
||||
)
|
||||
if not self.ibot_separate_head:
|
||||
inputs_for_student_head_list.append(buffer_tensor_patch_tokens.unsqueeze(0))
|
||||
else:
|
||||
student_global_masked_patch_tokens_after_head = self.student.ibot_head(buffer_tensor_patch_tokens)[
|
||||
:n_masked_patches
|
||||
]
|
||||
|
||||
# 2: run
|
||||
_attn_bias, cat_inputs = fmha.BlockDiagonalMask.from_tensor_list(inputs_for_student_head_list)
|
||||
outputs_list = _attn_bias.split(self.student.dino_head(cat_inputs))
|
||||
|
||||
# 3a: local crops cls tokens
|
||||
student_local_cls_tokens_after_head = outputs_list.pop(0).squeeze(0)
|
||||
|
||||
# 3b: global crops cls tokens
|
||||
student_global_cls_tokens_after_head = outputs_list.pop(0).squeeze(0)
|
||||
|
||||
# 3c: global crops patch tokens
|
||||
if do_ibot and not self.ibot_separate_head:
|
||||
student_global_masked_patch_tokens_after_head = outputs_list.pop(0).squeeze(0)[:n_masked_patches]
|
||||
|
||||
if n_local_crops > 0:
|
||||
dino_local_crops_loss = self.dino_loss(
|
||||
student_output_list=student_local_cls_tokens_after_head.chunk(n_local_crops),
|
||||
teacher_out_softmaxed_centered_list=teacher_dino_softmaxed_centered_list,
|
||||
) / (n_global_crops_loss_terms + n_local_crops_loss_terms)
|
||||
|
||||
# store for display
|
||||
loss_dict["dino_local_crops_loss"] = dino_local_crops_loss
|
||||
|
||||
# accumulate loss
|
||||
loss_accumulator += self.dino_loss_weight * dino_local_crops_loss
|
||||
|
||||
# process global crops
|
||||
loss_scales = 2 # this is here since we process global crops together
|
||||
|
||||
if do_dino:
|
||||
# compute loss
|
||||
dino_global_crops_loss = (
|
||||
self.dino_loss(
|
||||
student_output_list=[student_global_cls_tokens_after_head],
|
||||
teacher_out_softmaxed_centered_list=[
|
||||
teacher_dino_softmaxed_centered_list.flatten(0, 1)
|
||||
], # these were chunked and stacked in reverse so A is matched to B
|
||||
)
|
||||
* loss_scales
|
||||
/ (n_global_crops_loss_terms + n_local_crops_loss_terms)
|
||||
)
|
||||
|
||||
loss_dict["dino_global_crops_loss"] = dino_global_crops_loss
|
||||
|
||||
# accumulate loss
|
||||
loss_accumulator += self.dino_loss_weight * dino_global_crops_loss
|
||||
|
||||
student_cls_tokens = student_global_cls_tokens
|
||||
|
||||
if self.do_koleo:
|
||||
koleo_loss = self.cfg.dino.koleo_loss_weight * sum(
|
||||
self.koleo_loss(p) for p in student_cls_tokens.chunk(2)
|
||||
) # we don't apply koleo loss between cls tokens of a same image
|
||||
loss_accumulator += koleo_loss
|
||||
loss_dict["koleo_loss"] = (
|
||||
koleo_loss / loss_scales
|
||||
) # this is to display the same losses as before but we can remove eventually
|
||||
|
||||
if do_ibot:
|
||||
# compute loss
|
||||
ibot_patch_loss = (
|
||||
self.ibot_patch_loss.forward_masked(
|
||||
student_global_masked_patch_tokens_after_head,
|
||||
masked_teacher_ibot_softmaxed_centered,
|
||||
student_masks_flat=masks,
|
||||
n_masked_patches=n_masked_patches,
|
||||
masks_weight=masks_weight,
|
||||
)
|
||||
* loss_scales
|
||||
* ibot_loss_scale
|
||||
)
|
||||
|
||||
# store for display
|
||||
loss_dict["ibot_loss"] = ibot_patch_loss / 2
|
||||
|
||||
# accumulate loss
|
||||
loss_accumulator += self.ibot_loss_weight * ibot_patch_loss
|
||||
|
||||
self.backprop_loss(loss_accumulator)
|
||||
|
||||
self.fsdp_synchronize_streams()
|
||||
|
||||
return loss_dict
|
||||
|
||||
def fsdp_synchronize_streams(self):
|
||||
if self.need_to_synchronize_fsdp_streams:
|
||||
torch.cuda.synchronize()
|
||||
self.student.dino_head._streams = (
|
||||
self.teacher.dino_head._streams
|
||||
) = self.student.backbone._streams = self.teacher.backbone._streams
|
||||
self.need_to_synchronize_fsdp_streams = False
|
||||
|
||||
def update_teacher(self, m):
|
||||
student_param_list = []
|
||||
teacher_param_list = []
|
||||
with torch.no_grad():
|
||||
for k in self.student.keys():
|
||||
for ms, mt in zip(get_fsdp_modules(self.student[k]), get_fsdp_modules(self.teacher[k])):
|
||||
student_param_list += ms.params
|
||||
teacher_param_list += mt.params
|
||||
torch._foreach_mul_(teacher_param_list, m)
|
||||
torch._foreach_add_(teacher_param_list, student_param_list, alpha=1 - m)
|
||||
|
||||
def train(self):
|
||||
super().train()
|
||||
self.teacher.eval()
|
||||
|
||||
def get_maybe_fused_params_for_submodel(self, m):
|
||||
params_groups = get_params_groups_with_decay(
|
||||
model=m,
|
||||
lr_decay_rate=self.cfg.optim.layerwise_decay,
|
||||
patch_embed_lr_mult=self.cfg.optim.patch_embed_lr_mult,
|
||||
)
|
||||
fused_params_groups = fuse_params_groups(params_groups)
|
||||
logger.info("fusing param groups")
|
||||
|
||||
for g in fused_params_groups:
|
||||
g["foreach"] = True
|
||||
return fused_params_groups
|
||||
|
||||
def get_params_groups(self):
|
||||
all_params_groups = []
|
||||
for m in self.student.values():
|
||||
all_params_groups += self.get_maybe_fused_params_for_submodel(m)
|
||||
return all_params_groups
|
||||
|
||||
def prepare_for_distributed_training(self):
|
||||
logger.info("DISTRIBUTED FSDP -- preparing model for distributed training")
|
||||
if has_batchnorms(self.student):
|
||||
raise NotImplementedError
|
||||
# below will synchronize all student subnetworks across gpus:
|
||||
for k, v in self.student.items():
|
||||
self.teacher[k].load_state_dict(self.student[k].state_dict())
|
||||
student_model_cfg = self.cfg.compute_precision.student[k]
|
||||
self.student[k] = get_fsdp_wrapper(student_model_cfg, modules_to_wrap={BlockChunk})(self.student[k])
|
||||
teacher_model_cfg = self.cfg.compute_precision.teacher[k]
|
||||
self.teacher[k] = get_fsdp_wrapper(teacher_model_cfg, modules_to_wrap={BlockChunk})(self.teacher[k])
|
||||
@@ -0,0 +1,319 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from functools import partial
|
||||
|
||||
from fvcore.common.checkpoint import PeriodicCheckpointer
|
||||
import torch
|
||||
|
||||
from dinov2.data import SamplerType, make_data_loader, make_dataset
|
||||
from dinov2.data import collate_data_and_cast, DataAugmentationDINO, MaskingGenerator
|
||||
import dinov2.distributed as distributed
|
||||
from dinov2.fsdp import FSDPCheckpointer
|
||||
from dinov2.logging import MetricLogger
|
||||
from dinov2.utils.config import setup
|
||||
from dinov2.utils.utils import CosineScheduler
|
||||
|
||||
from dinov2.train.ssl_meta_arch import SSLMetaArch
|
||||
|
||||
|
||||
torch.backends.cuda.matmul.allow_tf32 = True # PyTorch 1.12 sets this to False by default
|
||||
logger = logging.getLogger("dinov2")
|
||||
|
||||
|
||||
def get_args_parser(add_help: bool = True):
|
||||
parser = argparse.ArgumentParser("DINOv2 training", add_help=add_help)
|
||||
parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
|
||||
parser.add_argument(
|
||||
"--no-resume",
|
||||
action="store_true",
|
||||
help="Whether to not attempt to resume from the checkpoint directory. ",
|
||||
)
|
||||
parser.add_argument("--eval-only", action="store_true", help="perform evaluation only")
|
||||
parser.add_argument("--eval", type=str, default="", help="Eval type to perform")
|
||||
parser.add_argument(
|
||||
"opts",
|
||||
help="""
|
||||
Modify config options at the end of the command. For Yacs configs, use
|
||||
space-separated "PATH.KEY VALUE" pairs.
|
||||
For python-based LazyConfig, use "path.key=value".
|
||||
""".strip(),
|
||||
default=None,
|
||||
nargs=argparse.REMAINDER,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
"--output_dir",
|
||||
default="",
|
||||
type=str,
|
||||
help="Output directory to save logs and checkpoints",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def build_optimizer(cfg, params_groups):
|
||||
return torch.optim.AdamW(params_groups, betas=(cfg.optim.adamw_beta1, cfg.optim.adamw_beta2))
|
||||
|
||||
|
||||
def build_schedulers(cfg):
|
||||
OFFICIAL_EPOCH_LENGTH = cfg.train.OFFICIAL_EPOCH_LENGTH
|
||||
lr = dict(
|
||||
base_value=cfg.optim["lr"],
|
||||
final_value=cfg.optim["min_lr"],
|
||||
total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH,
|
||||
warmup_iters=cfg.optim["warmup_epochs"] * OFFICIAL_EPOCH_LENGTH,
|
||||
start_warmup_value=0,
|
||||
)
|
||||
wd = dict(
|
||||
base_value=cfg.optim["weight_decay"],
|
||||
final_value=cfg.optim["weight_decay_end"],
|
||||
total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH,
|
||||
)
|
||||
momentum = dict(
|
||||
base_value=cfg.teacher["momentum_teacher"],
|
||||
final_value=cfg.teacher["final_momentum_teacher"],
|
||||
total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH,
|
||||
)
|
||||
teacher_temp = dict(
|
||||
base_value=cfg.teacher["teacher_temp"],
|
||||
final_value=cfg.teacher["teacher_temp"],
|
||||
total_iters=cfg.teacher["warmup_teacher_temp_epochs"] * OFFICIAL_EPOCH_LENGTH,
|
||||
warmup_iters=cfg.teacher["warmup_teacher_temp_epochs"] * OFFICIAL_EPOCH_LENGTH,
|
||||
start_warmup_value=cfg.teacher["warmup_teacher_temp"],
|
||||
)
|
||||
|
||||
lr_schedule = CosineScheduler(**lr)
|
||||
wd_schedule = CosineScheduler(**wd)
|
||||
momentum_schedule = CosineScheduler(**momentum)
|
||||
teacher_temp_schedule = CosineScheduler(**teacher_temp)
|
||||
last_layer_lr_schedule = CosineScheduler(**lr)
|
||||
|
||||
last_layer_lr_schedule.schedule[
|
||||
: cfg.optim["freeze_last_layer_epochs"] * OFFICIAL_EPOCH_LENGTH
|
||||
] = 0 # mimicking the original schedules
|
||||
|
||||
logger.info("Schedulers ready.")
|
||||
|
||||
return (
|
||||
lr_schedule,
|
||||
wd_schedule,
|
||||
momentum_schedule,
|
||||
teacher_temp_schedule,
|
||||
last_layer_lr_schedule,
|
||||
)
|
||||
|
||||
|
||||
def apply_optim_scheduler(optimizer, lr, wd, last_layer_lr):
|
||||
for param_group in optimizer.param_groups:
|
||||
is_last_layer = param_group["is_last_layer"]
|
||||
lr_multiplier = param_group["lr_multiplier"]
|
||||
wd_multiplier = param_group["wd_multiplier"]
|
||||
param_group["weight_decay"] = wd * wd_multiplier
|
||||
param_group["lr"] = (last_layer_lr if is_last_layer else lr) * lr_multiplier
|
||||
|
||||
|
||||
def do_test(cfg, model, iteration):
|
||||
new_state_dict = model.teacher.state_dict()
|
||||
|
||||
if distributed.is_main_process():
|
||||
iterstring = str(iteration)
|
||||
eval_dir = os.path.join(cfg.train.output_dir, "eval", iterstring)
|
||||
os.makedirs(eval_dir, exist_ok=True)
|
||||
# save teacher checkpoint
|
||||
teacher_ckp_path = os.path.join(eval_dir, "teacher_checkpoint.pth")
|
||||
torch.save({"teacher": new_state_dict}, teacher_ckp_path)
|
||||
|
||||
|
||||
def do_train(cfg, model, resume=False):
|
||||
model.train()
|
||||
inputs_dtype = torch.half
|
||||
fp16_scaler = model.fp16_scaler # for mixed precision training
|
||||
|
||||
# setup optimizer
|
||||
|
||||
optimizer = build_optimizer(cfg, model.get_params_groups())
|
||||
(
|
||||
lr_schedule,
|
||||
wd_schedule,
|
||||
momentum_schedule,
|
||||
teacher_temp_schedule,
|
||||
last_layer_lr_schedule,
|
||||
) = build_schedulers(cfg)
|
||||
|
||||
# checkpointer
|
||||
checkpointer = FSDPCheckpointer(model, cfg.train.output_dir, optimizer=optimizer, save_to_disk=True)
|
||||
|
||||
start_iter = checkpointer.resume_or_load(cfg.MODEL.WEIGHTS, resume=resume).get("iteration", -1) + 1
|
||||
|
||||
OFFICIAL_EPOCH_LENGTH = cfg.train.OFFICIAL_EPOCH_LENGTH
|
||||
max_iter = cfg.optim.epochs * OFFICIAL_EPOCH_LENGTH
|
||||
|
||||
periodic_checkpointer = PeriodicCheckpointer(
|
||||
checkpointer,
|
||||
period=3 * OFFICIAL_EPOCH_LENGTH,
|
||||
max_iter=max_iter,
|
||||
max_to_keep=3,
|
||||
)
|
||||
|
||||
# setup data preprocessing
|
||||
|
||||
img_size = cfg.crops.global_crops_size
|
||||
patch_size = cfg.student.patch_size
|
||||
n_tokens = (img_size // patch_size) ** 2
|
||||
mask_generator = MaskingGenerator(
|
||||
input_size=(img_size // patch_size, img_size // patch_size),
|
||||
max_num_patches=0.5 * img_size // patch_size * img_size // patch_size,
|
||||
)
|
||||
|
||||
data_transform = DataAugmentationDINO(
|
||||
cfg.crops.global_crops_scale,
|
||||
cfg.crops.local_crops_scale,
|
||||
cfg.crops.local_crops_number,
|
||||
global_crops_size=cfg.crops.global_crops_size,
|
||||
local_crops_size=cfg.crops.local_crops_size,
|
||||
)
|
||||
|
||||
collate_fn = partial(
|
||||
collate_data_and_cast,
|
||||
mask_ratio_tuple=cfg.ibot.mask_ratio_min_max,
|
||||
mask_probability=cfg.ibot.mask_sample_probability,
|
||||
n_tokens=n_tokens,
|
||||
mask_generator=mask_generator,
|
||||
dtype=inputs_dtype,
|
||||
)
|
||||
|
||||
# setup data loader
|
||||
|
||||
dataset = make_dataset(
|
||||
dataset_str=cfg.train.dataset_path,
|
||||
transform=data_transform,
|
||||
target_transform=lambda _: (),
|
||||
)
|
||||
# sampler_type = SamplerType.INFINITE
|
||||
sampler_type = SamplerType.SHARDED_INFINITE
|
||||
data_loader = make_data_loader(
|
||||
dataset=dataset,
|
||||
batch_size=cfg.train.batch_size_per_gpu,
|
||||
num_workers=cfg.train.num_workers,
|
||||
shuffle=True,
|
||||
seed=start_iter, # TODO: Fix this -- cfg.train.seed
|
||||
sampler_type=sampler_type,
|
||||
sampler_advance=0, # TODO(qas): fix this -- start_iter * cfg.train.batch_size_per_gpu,
|
||||
drop_last=True,
|
||||
collate_fn=collate_fn,
|
||||
)
|
||||
|
||||
# training loop
|
||||
|
||||
iteration = start_iter
|
||||
|
||||
logger.info("Starting training from iteration {}".format(start_iter))
|
||||
metrics_file = os.path.join(cfg.train.output_dir, "training_metrics.json")
|
||||
metric_logger = MetricLogger(delimiter=" ", output_file=metrics_file)
|
||||
header = "Training"
|
||||
|
||||
for data in metric_logger.log_every(
|
||||
data_loader,
|
||||
10,
|
||||
header,
|
||||
max_iter,
|
||||
start_iter,
|
||||
):
|
||||
current_batch_size = data["collated_global_crops"].shape[0] / 2
|
||||
if iteration > max_iter:
|
||||
return
|
||||
|
||||
# apply schedules
|
||||
|
||||
lr = lr_schedule[iteration]
|
||||
wd = wd_schedule[iteration]
|
||||
mom = momentum_schedule[iteration]
|
||||
teacher_temp = teacher_temp_schedule[iteration]
|
||||
last_layer_lr = last_layer_lr_schedule[iteration]
|
||||
apply_optim_scheduler(optimizer, lr, wd, last_layer_lr)
|
||||
|
||||
# compute losses
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss_dict = model.forward_backward(data, teacher_temp=teacher_temp)
|
||||
|
||||
# clip gradients
|
||||
|
||||
if fp16_scaler is not None:
|
||||
if cfg.optim.clip_grad:
|
||||
fp16_scaler.unscale_(optimizer)
|
||||
for v in model.student.values():
|
||||
v.clip_grad_norm_(cfg.optim.clip_grad)
|
||||
fp16_scaler.step(optimizer)
|
||||
fp16_scaler.update()
|
||||
else:
|
||||
if cfg.optim.clip_grad:
|
||||
for v in model.student.values():
|
||||
v.clip_grad_norm_(cfg.optim.clip_grad)
|
||||
optimizer.step()
|
||||
|
||||
# perform teacher EMA update
|
||||
|
||||
model.update_teacher(mom)
|
||||
|
||||
# logging
|
||||
|
||||
if distributed.get_global_size() > 1:
|
||||
for v in loss_dict.values():
|
||||
torch.distributed.all_reduce(v)
|
||||
loss_dict_reduced = {k: v.item() / distributed.get_global_size() for k, v in loss_dict.items()}
|
||||
|
||||
if math.isnan(sum(loss_dict_reduced.values())):
|
||||
logger.info("NaN detected")
|
||||
raise AssertionError
|
||||
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
|
||||
|
||||
metric_logger.update(lr=lr)
|
||||
metric_logger.update(wd=wd)
|
||||
metric_logger.update(mom=mom)
|
||||
metric_logger.update(last_layer_lr=last_layer_lr)
|
||||
metric_logger.update(current_batch_size=current_batch_size)
|
||||
metric_logger.update(total_loss=losses_reduced, **loss_dict_reduced)
|
||||
|
||||
# checkpointing and testing
|
||||
|
||||
if cfg.evaluation.eval_period_iterations > 0 and (iteration + 1) % cfg.evaluation.eval_period_iterations == 0:
|
||||
do_test(cfg, model, f"training_{iteration}")
|
||||
torch.cuda.synchronize()
|
||||
periodic_checkpointer.step(iteration)
|
||||
|
||||
iteration = iteration + 1
|
||||
metric_logger.synchronize_between_processes()
|
||||
return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
|
||||
|
||||
|
||||
def main(args):
|
||||
cfg = setup(args)
|
||||
|
||||
model = SSLMetaArch(cfg).to(torch.device("cuda"))
|
||||
model.prepare_for_distributed_training()
|
||||
|
||||
logger.info("Model:\n{}".format(model))
|
||||
if args.eval_only:
|
||||
iteration = (
|
||||
FSDPCheckpointer(model, save_dir=cfg.train.output_dir)
|
||||
.resume_or_load(cfg.MODEL.WEIGHTS, resume=not args.no_resume)
|
||||
.get("iteration", -1)
|
||||
+ 1
|
||||
)
|
||||
return do_test(cfg, model, f"manual_{iteration}")
|
||||
|
||||
do_train(cfg, model, resume=not args.no_resume)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = get_args_parser(add_help=True).parse_args()
|
||||
main(args)
|
||||
Reference in New Issue
Block a user