Spaces:
Sleeping
Sleeping
# Copyright (c) Meta Platforms, Inc. and affiliates. | |
# | |
# This source code is licensed under the Apache License, Version 2.0 | |
# 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) | |