diff --git a/dino/__init__.py b/dino/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/eval_copy_detection.py b/dino/eval_copy_detection.py similarity index 94% rename from eval_copy_detection.py rename to dino/eval_copy_detection.py index 73dcd5078..b8c2943a4 100644 --- a/eval_copy_detection.py +++ b/dino/eval_copy_detection.py @@ -25,9 +25,9 @@ from PIL import Image, ImageFile import numpy as np -import utils -import vision_transformer as vits -from eval_knn import extract_features +import dino.utils +import dino.vision_transformer as vits +from dino.eval_knn import extract_features class CopydaysDataset(): @@ -161,7 +161,7 @@ def extract_features(image_list, model, args): num_workers=args.num_workers, drop_last=False, sampler=torch.utils.data.DistributedSampler(tempdataset, shuffle=False)) features = None - for samples, index in utils.MetricLogger(delimiter=" ").log_every(data_loader, 10): + for samples, index in dino.utils.MetricLogger(delimiter=" ").log_every(data_loader, 10): samples, index = samples.cuda(non_blocking=True), index.cuda(non_blocking=True) feats = model.get_intermediate_layers(samples, n=1)[0].clone() @@ -215,7 +215,7 @@ def extract_features(image_list, model, args): parser.add_argument('--imsize', default=320, type=int, help='Image size (square image)') parser.add_argument('--batch_size_per_gpu', default=16, type=int, help='Per-GPU batch-size') parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.") - parser.add_argument('--use_cuda', default=True, type=utils.bool_flag) + parser.add_argument('--use_cuda', default=True, type=dino.utils.bool_flag) parser.add_argument('--arch', default='vit_base', type=str, help='Architecture') parser.add_argument('--patch_size', default=8, type=int, help='Patch resolution of the model.') parser.add_argument("--checkpoint_key", default="teacher", type=str, @@ -226,8 +226,8 @@ def extract_features(image_list, model, args): parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.") args = parser.parse_args() - utils.init_distributed_mode(args) - print("git:\n {}\n".format(utils.get_sha())) + dino.utils.init_distributed_mode(args) + print("git:\n {}\n".format(dino.utils.get_sha())) print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))) cudnn.benchmark = True @@ -241,7 +241,7 @@ def extract_features(image_list, model, args): if args.use_cuda: model.cuda() model.eval() - utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size) + dino.utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size) dataset = CopydaysDataset(args.data_path) @@ -250,7 +250,7 @@ def extract_features(image_list, model, args): queries = [] for q in dataset.query_blocks: queries.append(extract_features(dataset.get_block(q), model, args)) - if utils.get_rank() == 0: + if dino.utils.get_rank() == 0: queries = torch.cat(queries) print(f"Extraction of queries features done. Shape: {queries.shape}") @@ -264,7 +264,7 @@ def extract_features(image_list, model, args): print("Using distractors...") list_distractors = [os.path.join(args.distractors_path, s) for s in os.listdir(args.distractors_path) if is_image_file(s)] database.append(extract_features(list_distractors, model, args)) - if utils.get_rank() == 0: + if dino.utils.get_rank() == 0: database = torch.cat(database) print(f"Extraction of database and distractors features done. Shape: {database.shape}") @@ -273,12 +273,12 @@ def extract_features(image_list, model, args): print(f"Extracting features on images from {args.whitening_path} for learning the whitening operator.") list_whit = [os.path.join(args.whitening_path, s) for s in os.listdir(args.whitening_path) if is_image_file(s)] features_for_whitening = extract_features(list_whit, model, args) - if utils.get_rank() == 0: + if dino.utils.get_rank() == 0: # center mean_feature = torch.mean(features_for_whitening, dim=0) database -= mean_feature queries -= mean_feature - pca = utils.PCA(dim=database.shape[-1], whit=0.5) + pca = dino.utils.PCA(dim=database.shape[-1], whit=0.5) # compute covariance cov = torch.mm(features_for_whitening.T, features_for_whitening) / features_for_whitening.shape[0] pca.train_pca(cov.cpu().numpy()) @@ -286,7 +286,7 @@ def extract_features(image_list, model, args): queries = pca.apply(queries) # ============ Copy detection ... ============ - if utils.get_rank() == 0: + if dino.utils.get_rank() == 0: # l2 normalize the features database = nn.functional.normalize(database, dim=1, p=2) queries = nn.functional.normalize(queries, dim=1, p=2) diff --git a/eval_image_retrieval.py b/dino/eval_image_retrieval.py similarity index 93% rename from eval_image_retrieval.py rename to dino/eval_image_retrieval.py index 999f8c900..36c55c841 100644 --- a/eval_image_retrieval.py +++ b/dino/eval_image_retrieval.py @@ -25,9 +25,9 @@ from PIL import Image, ImageFile import numpy as np -import utils -import vision_transformer as vits -from eval_knn import extract_features +import dino.utils +import dino.vision_transformer as vits +from dino.eval_knn import extract_features class OxfordParisDataset(torch.utils.data.Dataset): @@ -83,10 +83,10 @@ def config_qimname(cfg, i): parser = argparse.ArgumentParser('Image Retrieval on revisited Paris and Oxford') parser.add_argument('--data_path', default='/path/to/revisited_paris_oxford/', type=str) parser.add_argument('--dataset', default='roxford5k', type=str, choices=['roxford5k', 'rparis6k']) - parser.add_argument('--multiscale', default=False, type=utils.bool_flag) + parser.add_argument('--multiscale', default=False, type=dino.utils.bool_flag) parser.add_argument('--imsize', default=224, type=int, help='Image size') parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.") - parser.add_argument('--use_cuda', default=True, type=utils.bool_flag) + parser.add_argument('--use_cuda', default=True, type=dino.utils.bool_flag) parser.add_argument('--arch', default='vit_small', type=str, help='Architecture') parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.') parser.add_argument("--checkpoint_key", default="teacher", type=str, @@ -97,8 +97,8 @@ def config_qimname(cfg, i): parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.") args = parser.parse_args() - utils.init_distributed_mode(args) - print("git:\n {}\n".format(utils.get_sha())) + dino.utils.init_distributed_mode(args) + print("git:\n {}\n".format(dino.utils.get_sha())) print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))) cudnn.benchmark = True @@ -165,7 +165,7 @@ def config_qimname(cfg, i): train_features = extract_features(model, data_loader_train, args.use_cuda, multiscale=args.multiscale) query_features = extract_features(model, data_loader_query, args.use_cuda, multiscale=args.multiscale) - if utils.get_rank() == 0: # only rank 0 will work from now on + if dino.utils.get_rank() == 0: # only rank 0 will work from now on # normalize features train_features = nn.functional.normalize(train_features, dim=1, p=2) query_features = nn.functional.normalize(query_features, dim=1, p=2) @@ -187,7 +187,7 @@ def config_qimname(cfg, i): g['ok'] = np.concatenate([gnd[i]['easy'], gnd[i]['hard']]) g['junk'] = np.concatenate([gnd[i]['junk']]) gnd_t.append(g) - mapM, apsM, mprM, prsM = utils.compute_map(ranks, gnd_t, ks) + mapM, apsM, mprM, prsM = dino.utils.compute_map(ranks, gnd_t, ks) # search for hard gnd_t = [] for i in range(len(gnd)): @@ -195,7 +195,7 @@ def config_qimname(cfg, i): g['ok'] = np.concatenate([gnd[i]['hard']]) g['junk'] = np.concatenate([gnd[i]['junk'], gnd[i]['easy']]) gnd_t.append(g) - mapH, apsH, mprH, prsH = utils.compute_map(ranks, gnd_t, ks) + mapH, apsH, mprH, prsH = dino.utils.compute_map(ranks, gnd_t, ks) print('>> {}: mAP M: {}, H: {}'.format(args.dataset, np.around(mapM*100, decimals=2), np.around(mapH*100, decimals=2))) print('>> {}: mP@k{} M: {}, H: {}'.format(args.dataset, np.array(ks), np.around(mprM*100, decimals=2), np.around(mprH*100, decimals=2))) dist.barrier() diff --git a/eval_knn.py b/dino/eval_knn.py similarity index 95% rename from eval_knn.py rename to dino/eval_knn.py index fe99a2604..de4fec766 100644 --- a/eval_knn.py +++ b/dino/eval_knn.py @@ -23,8 +23,8 @@ from torchvision import transforms as pth_transforms from torchvision import models as torchvision_models -import utils -import vision_transformer as vits +import dino.utils +import dino.vision_transformer as vits def extract_feature_pipeline(args): @@ -68,7 +68,7 @@ def extract_feature_pipeline(args): print(f"Architecture {args.arch} non supported") sys.exit(1) model.cuda() - utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size) + dino.utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size) model.eval() # ============ extract features ... ============ @@ -77,7 +77,7 @@ def extract_feature_pipeline(args): print("Extracting features for val set...") test_features = extract_features(model, data_loader_val, args.use_cuda) - if utils.get_rank() == 0: + if dino.utils.get_rank() == 0: train_features = nn.functional.normalize(train_features, dim=1, p=2) test_features = nn.functional.normalize(test_features, dim=1, p=2) @@ -94,13 +94,13 @@ def extract_feature_pipeline(args): @torch.no_grad() def extract_features(model, data_loader, use_cuda=True, multiscale=False): - metric_logger = utils.MetricLogger(delimiter=" ") + metric_logger = dino.utils.MetricLogger(delimiter=" ") features = None for samples, index in metric_logger.log_every(data_loader, 10): samples = samples.cuda(non_blocking=True) index = index.cuda(non_blocking=True) if multiscale: - feats = utils.multi_scale(samples, model) + feats = dino.utils.multi_scale(samples, model) else: feats = model(samples).clone() @@ -196,7 +196,7 @@ def __getitem__(self, idx): parser.add_argument('--temperature', default=0.07, type=float, help='Temperature used in the voting coefficient') parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.") - parser.add_argument('--use_cuda', default=True, type=utils.bool_flag, + parser.add_argument('--use_cuda', default=True, type=dino.utils.bool_flag, help="Should we store the features on GPU? We recommend setting this to False if you encounter OOM") parser.add_argument('--arch', default='vit_small', type=str, help='Architecture') parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.') @@ -213,8 +213,8 @@ def __getitem__(self, idx): parser.add_argument('--data_path', default='/path/to/imagenet/', type=str) args = parser.parse_args() - utils.init_distributed_mode(args) - print("git:\n {}\n".format(utils.get_sha())) + dino.utils.init_distributed_mode(args) + print("git:\n {}\n".format(dino.utils.get_sha())) print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))) cudnn.benchmark = True @@ -227,7 +227,7 @@ def __getitem__(self, idx): # need to extract features ! train_features, test_features, train_labels, test_labels = extract_feature_pipeline(args) - if utils.get_rank() == 0: + if dino.utils.get_rank() == 0: if args.use_cuda: train_features = train_features.cuda() test_features = test_features.cuda() diff --git a/eval_linear.py b/dino/eval_linear.py similarity index 92% rename from eval_linear.py rename to dino/eval_linear.py index cdef16b47..aee13d3b5 100644 --- a/eval_linear.py +++ b/dino/eval_linear.py @@ -24,13 +24,13 @@ from torchvision import transforms as pth_transforms from torchvision import models as torchvision_models -import utils -import vision_transformer as vits +import dino.utils +import dino.vision_transformer as vits def eval_linear(args): - utils.init_distributed_mode(args) - print("git:\n {}\n".format(utils.get_sha())) + dino.utils.init_distributed_mode(args) + print("git:\n {}\n".format(dino.utils.get_sha())) print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))) cudnn.benchmark = True @@ -54,7 +54,7 @@ def eval_linear(args): model.cuda() model.eval() # load weights to evaluate - utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size) + dino.utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size) print(f"Model {args.arch} built.") linear_classifier = LinearClassifier(embed_dim, num_labels=args.num_labels) @@ -77,7 +77,7 @@ def eval_linear(args): ) if args.evaluate: - utils.load_pretrained_linear_weights(linear_classifier, args.arch, args.patch_size) + dino.utils.load_pretrained_linear_weights(linear_classifier, args.arch, args.patch_size) test_stats = validate_network(val_loader, model, linear_classifier, args.n_last_blocks, args.avgpool_patchtokens) print(f"Accuracy of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%") return @@ -102,7 +102,7 @@ def eval_linear(args): # set optimizer optimizer = torch.optim.SGD( linear_classifier.parameters(), - args.lr * (args.batch_size_per_gpu * utils.get_world_size()) / 256., # linear scaling rule + args.lr * (args.batch_size_per_gpu * dino.utils.get_world_size()) / 256., # linear scaling rule momentum=0.9, weight_decay=0, # we do not apply weight decay ) @@ -110,7 +110,7 @@ def eval_linear(args): # Optionally resume from a checkpoint to_restore = {"epoch": 0, "best_acc": 0.} - utils.restart_from_checkpoint( + dino.utils.restart_from_checkpoint( os.path.join(args.output_dir, "checkpoint.pth.tar"), run_variables=to_restore, state_dict=linear_classifier, @@ -135,7 +135,7 @@ def eval_linear(args): print(f'Max accuracy so far: {best_acc:.2f}%') log_stats = {**{k: v for k, v in log_stats.items()}, **{f'test_{k}': v for k, v in test_stats.items()}} - if utils.is_main_process(): + if dino.utils.is_main_process(): with (Path(args.output_dir) / "log.txt").open("a") as f: f.write(json.dumps(log_stats) + "\n") save_dict = { @@ -152,8 +152,8 @@ def eval_linear(args): def train(model, linear_classifier, optimizer, loader, epoch, n, avgpool): linear_classifier.train() - metric_logger = utils.MetricLogger(delimiter=" ") - metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}')) + metric_logger = dino.utils.MetricLogger(delimiter=" ") + metric_logger.add_meter('lr', dino.utils.SmoothedValue(window_size=1, fmt='{value:.6f}')) header = 'Epoch: [{}]'.format(epoch) for (inp, target) in metric_logger.log_every(loader, 20, header): # move to gpu @@ -195,7 +195,7 @@ def train(model, linear_classifier, optimizer, loader, epoch, n, avgpool): @torch.no_grad() def validate_network(val_loader, model, linear_classifier, n, avgpool): linear_classifier.eval() - metric_logger = utils.MetricLogger(delimiter=" ") + metric_logger = dino.utils.MetricLogger(delimiter=" ") header = 'Test:' for inp, target in metric_logger.log_every(val_loader, 20, header): # move to gpu @@ -216,9 +216,9 @@ def validate_network(val_loader, model, linear_classifier, n, avgpool): loss = nn.CrossEntropyLoss()(output, target) if linear_classifier.module.num_labels >= 5: - acc1, acc5 = utils.accuracy(output, target, topk=(1, 5)) + acc1, acc5 = dino.utils.accuracy(output, target, topk=(1, 5)) else: - acc1, = utils.accuracy(output, target, topk=(1,)) + acc1, = dino.utils.accuracy(output, target, topk=(1,)) batch_size = inp.shape[0] metric_logger.update(loss=loss.item()) @@ -255,7 +255,7 @@ def forward(self, x): parser = argparse.ArgumentParser('Evaluation with linear classification on ImageNet') parser.add_argument('--n_last_blocks', default=4, type=int, help="""Concatenate [CLS] tokens for the `n` last blocks. We use `n=4` when evaluating ViT-Small and `n=1` with ViT-Base.""") - parser.add_argument('--avgpool_patchtokens', default=False, type=utils.bool_flag, + parser.add_argument('--avgpool_patchtokens', default=False, type=dino.utils.bool_flag, help="""Whether ot not to concatenate the global average pooled features to the [CLS] token. We typically set this to False for ViT-Small and to True with ViT-Base.""") parser.add_argument('--arch', default='vit_small', type=str, help='Architecture') diff --git a/eval_video_segmentation.py b/dino/eval_video_segmentation.py similarity index 98% rename from eval_video_segmentation.py rename to dino/eval_video_segmentation.py index 08a18c475..1bef4ffce 100644 --- a/eval_video_segmentation.py +++ b/dino/eval_video_segmentation.py @@ -30,8 +30,8 @@ from PIL import Image from torchvision import transforms -import utils -import vision_transformer as vits +import dino.utils +import dino.vision_transformer as vits @torch.no_grad() @@ -264,14 +264,14 @@ def color_normalize(x, mean=[0.485, 0.456, 0.406], std=[0.228, 0.224, 0.225]): parser.add_argument("--bs", type=int, default=6, help="Batch size, try to reduce if OOM") args = parser.parse_args() - print("git:\n {}\n".format(utils.get_sha())) + print("git:\n {}\n".format(dino.utils.get_sha())) print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))) # building network model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0) print(f"Model {args.arch} {args.patch_size}x{args.patch_size} built.") model.cuda() - utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size) + dino.utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size) for param in model.parameters(): param.requires_grad = False model.eval() diff --git a/hubconf.py b/dino/hubconf.py similarity index 99% rename from hubconf.py rename to dino/hubconf.py index 3709271ed..39ea0d2b6 100644 --- a/hubconf.py +++ b/dino/hubconf.py @@ -14,7 +14,7 @@ import torch from torchvision.models.resnet import resnet50 -import vision_transformer as vits +import dino.vision_transformer as vits dependencies = ["torch", "torchvision"] diff --git a/main_dino.py b/dino/main_dino.py similarity index 92% rename from main_dino.py rename to dino/main_dino.py index cade9873d..ff88e7fb1 100644 --- a/main_dino.py +++ b/dino/main_dino.py @@ -30,9 +30,9 @@ from torchvision import datasets, transforms from torchvision import models as torchvision_models -import utils -import vision_transformer as vits -from vision_transformer import DINOHead +import dino.utils +import dino.vision_transformer as vits +from dino.vision_transformer import DINOHead torchvision_archs = sorted(name for name in torchvision_models.__dict__ if name.islower() and not name.startswith("__") @@ -54,14 +54,14 @@ def get_args_parser(): mixed precision training (--use_fp16 false) to avoid unstabilities.""") parser.add_argument('--out_dim', default=65536, type=int, help="""Dimensionality of the DINO head output. For complex and large datasets large values (like 65k) work well.""") - parser.add_argument('--norm_last_layer', default=True, type=utils.bool_flag, + parser.add_argument('--norm_last_layer', default=True, type=dino.utils.bool_flag, help="""Whether or not to weight normalize the last layer of the DINO head. Not normalizing leads to better performance but can make the training unstable. In our experiments, we typically set this paramater to False with vit_small and True with vit_base.""") parser.add_argument('--momentum_teacher', default=0.996, type=float, help="""Base EMA parameter for teacher update. The value is increased to 1 during training with cosine schedule. We recommend setting a higher value with small batches: for example use 0.9995 with batch size of 256.""") - parser.add_argument('--use_bn_in_head', default=False, type=utils.bool_flag, + parser.add_argument('--use_bn_in_head', default=False, type=dino.utils.bool_flag, help="Whether to use batch normalizations in projection head (Default: False)") # Temperature teacher parameters @@ -75,7 +75,7 @@ def get_args_parser(): help='Number of warmup epochs for the teacher temperature (Default: 30).') # Training/Optimization parameters - parser.add_argument('--use_fp16', type=utils.bool_flag, default=True, help="""Whether or not + parser.add_argument('--use_fp16', type=dino.utils.bool_flag, default=True, help="""Whether or not to use half precision for training. Improves training time and memory requirements, but can provoke instability and slight decay of performance. We recommend disabling mixed precision if the loss is unstable, if reducing the patch size or if training with bigger ViTs.""") @@ -130,9 +130,9 @@ def get_args_parser(): def train_dino(args): - utils.init_distributed_mode(args) - utils.fix_random_seeds(args.seed) - print("git:\n {}\n".format(utils.get_sha())) + dino.utils.init_distributed_mode(args) + dino.utils.fix_random_seeds(args.seed) + print("git:\n {}\n".format(dino.utils.get_sha())) print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))) cudnn.benchmark = True @@ -180,20 +180,20 @@ def train_dino(args): print(f"Unknow architecture: {args.arch}") # multi-crop wrapper handles forward with inputs of different resolutions - student = utils.MultiCropWrapper(student, DINOHead( + student = dino.utils.MultiCropWrapper(student, DINOHead( embed_dim, args.out_dim, use_bn=args.use_bn_in_head, norm_last_layer=args.norm_last_layer, )) - teacher = utils.MultiCropWrapper( + teacher = dino.utils.MultiCropWrapper( teacher, DINOHead(embed_dim, args.out_dim, args.use_bn_in_head), ) # move networks to gpu student, teacher = student.cuda(), teacher.cuda() # synchronize batch norms (if any) - if utils.has_batchnorms(student): + if dino.utils.has_batchnorms(student): student = nn.SyncBatchNorm.convert_sync_batchnorm(student) teacher = nn.SyncBatchNorm.convert_sync_batchnorm(teacher) @@ -222,38 +222,38 @@ def train_dino(args): ).cuda() # ============ preparing optimizer ... ============ - params_groups = utils.get_params_groups(student) + params_groups = dino.utils.get_params_groups(student) if args.optimizer == "adamw": optimizer = torch.optim.AdamW(params_groups) # to use with ViTs elif args.optimizer == "sgd": optimizer = torch.optim.SGD(params_groups, lr=0, momentum=0.9) # lr is set by scheduler elif args.optimizer == "lars": - optimizer = utils.LARS(params_groups) # to use with convnet and large batches + optimizer = dino.utils.LARS(params_groups) # to use with convnet and large batches # for mixed precision training fp16_scaler = None if args.use_fp16: fp16_scaler = torch.cuda.amp.GradScaler() # ============ init schedulers ... ============ - lr_schedule = utils.cosine_scheduler( - args.lr * (args.batch_size_per_gpu * utils.get_world_size()) / 256., # linear scaling rule + lr_schedule = dino.utils.cosine_scheduler( + args.lr * (args.batch_size_per_gpu * dino.utils.get_world_size()) / 256., # linear scaling rule args.min_lr, args.epochs, len(data_loader), warmup_epochs=args.warmup_epochs, ) - wd_schedule = utils.cosine_scheduler( + wd_schedule = dino.utils.cosine_scheduler( args.weight_decay, args.weight_decay_end, args.epochs, len(data_loader), ) # momentum parameter is increased to 1. during training with a cosine schedule - momentum_schedule = utils.cosine_scheduler(args.momentum_teacher, 1, + momentum_schedule = dino.utils.cosine_scheduler(args.momentum_teacher, 1, args.epochs, len(data_loader)) print(f"Loss, optimizer and schedulers ready.") # ============ optionally resume training ... ============ to_restore = {"epoch": 0} - utils.restart_from_checkpoint( + dino.utils.restart_from_checkpoint( os.path.join(args.output_dir, "checkpoint.pth"), run_variables=to_restore, student=student, @@ -285,12 +285,12 @@ def train_dino(args): } if fp16_scaler is not None: save_dict['fp16_scaler'] = fp16_scaler.state_dict() - utils.save_on_master(save_dict, os.path.join(args.output_dir, 'checkpoint.pth')) + dino.utils.save_on_master(save_dict, os.path.join(args.output_dir, 'checkpoint.pth')) if args.saveckp_freq and epoch % args.saveckp_freq == 0: - utils.save_on_master(save_dict, os.path.join(args.output_dir, f'checkpoint{epoch:04}.pth')) + dino.utils.save_on_master(save_dict, os.path.join(args.output_dir, f'checkpoint{epoch:04}.pth')) log_stats = {**{f'train_{k}': v for k, v in train_stats.items()}, 'epoch': epoch} - if utils.is_main_process(): + if dino.utils.is_main_process(): with (Path(args.output_dir) / "log.txt").open("a") as f: f.write(json.dumps(log_stats) + "\n") total_time = time.time() - start_time @@ -301,7 +301,7 @@ def train_dino(args): def train_one_epoch(student, teacher, teacher_without_ddp, dino_loss, data_loader, optimizer, lr_schedule, wd_schedule, momentum_schedule,epoch, fp16_scaler, args): - metric_logger = utils.MetricLogger(delimiter=" ") + metric_logger = dino.utils.MetricLogger(delimiter=" ") header = 'Epoch: [{}/{}]'.format(epoch, args.epochs) for it, (images, _) in enumerate(metric_logger.log_every(data_loader, 10, header)): # update weight decay and learning rate according to their schedule @@ -329,16 +329,16 @@ def train_one_epoch(student, teacher, teacher_without_ddp, dino_loss, data_loade if fp16_scaler is None: loss.backward() if args.clip_grad: - param_norms = utils.clip_gradients(student, args.clip_grad) - utils.cancel_gradients_last_layer(epoch, student, + param_norms = dino.utils.clip_gradients(student, args.clip_grad) + dino.utils.cancel_gradients_last_layer(epoch, student, args.freeze_last_layer) optimizer.step() else: fp16_scaler.scale(loss).backward() if args.clip_grad: fp16_scaler.unscale_(optimizer) # unscale the gradients of optimizer's assigned params in-place - param_norms = utils.clip_gradients(student, args.clip_grad) - utils.cancel_gradients_last_layer(epoch, student, + param_norms = dino.utils.clip_gradients(student, args.clip_grad) + dino.utils.cancel_gradients_last_layer(epoch, student, args.freeze_last_layer) fp16_scaler.step(optimizer) fp16_scaler.update() @@ -435,15 +435,15 @@ def __init__(self, global_crops_scale, local_crops_scale, local_crops_number): self.global_transfo1 = transforms.Compose([ transforms.RandomResizedCrop(224, scale=global_crops_scale, interpolation=Image.BICUBIC), flip_and_color_jitter, - utils.GaussianBlur(1.0), + dino.utils.GaussianBlur(1.0), normalize, ]) # second global crop self.global_transfo2 = transforms.Compose([ transforms.RandomResizedCrop(224, scale=global_crops_scale, interpolation=Image.BICUBIC), flip_and_color_jitter, - utils.GaussianBlur(0.1), - utils.Solarization(0.2), + dino.utils.GaussianBlur(0.1), + dino.utils.Solarization(0.2), normalize, ]) # transformation for the local small crops @@ -451,7 +451,7 @@ def __init__(self, global_crops_scale, local_crops_scale, local_crops_number): self.local_transfo = transforms.Compose([ transforms.RandomResizedCrop(96, scale=local_crops_scale, interpolation=Image.BICUBIC), flip_and_color_jitter, - utils.GaussianBlur(p=0.5), + dino.utils.GaussianBlur(p=0.5), normalize, ]) diff --git a/run_with_submitit.py b/dino/run_with_submitit.py similarity index 96% rename from run_with_submitit.py rename to dino/run_with_submitit.py index 33d4116f2..374c1fde2 100644 --- a/run_with_submitit.py +++ b/dino/run_with_submitit.py @@ -20,12 +20,12 @@ import uuid from pathlib import Path -import main_dino +import dino.main_dino import submitit def parse_args(): - parser = argparse.ArgumentParser("Submitit for DINO", parents=[main_dino.get_args_parser()]) + parser = argparse.ArgumentParser("Submitit for DINO", parents=[dino.main_dino.get_args_parser()]) parser.add_argument("--ngpus", default=8, type=int, help="Number of gpus to request on each node") parser.add_argument("--nodes", default=2, type=int, help="Number of nodes to request") parser.add_argument("--timeout", default=2800, type=int, help="Duration of the job") @@ -60,10 +60,10 @@ def __init__(self, args): self.args = args def __call__(self): - import main_dino + import dino.main_dino self._setup_gpu_args() - main_dino.train_dino(self.args) + dino.main_dino.train_dino(self.args) def checkpoint(self): import os diff --git a/utils.py b/dino/utils.py similarity index 99% rename from utils.py rename to dino/utils.py index 958625012..13905b2e8 100644 --- a/utils.py +++ b/dino/utils.py @@ -103,6 +103,7 @@ def load_pretrained_weights(model, pretrained_weights, checkpoint_key, model_nam url = "dino_resnet50_pretrain/dino_resnet50_pretrain.pth" if url is not None: print("Since no pretrained weights have been provided, we load the reference pretrained DINO weights.") + print(f"Going to download {url=} via torch.hub") state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/" + url) model.load_state_dict(state_dict, strict=True) else: @@ -123,6 +124,7 @@ def load_pretrained_linear_weights(linear_classifier, model_name, patch_size): url = "dino_resnet50_pretrain/dino_resnet50_linearweights.pth" if url is not None: print("We load the reference pretrained linear weights.") + print(f"Going to download {url=} via torch.hub") state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/" + url)["state_dict"] linear_classifier.load_state_dict(state_dict, strict=True) else: diff --git a/video_generation.py b/dino/video_generation.py similarity index 99% rename from video_generation.py rename to dino/video_generation.py index 94da9836a..f90ae8664 100644 --- a/video_generation.py +++ b/dino/video_generation.py @@ -26,8 +26,8 @@ import numpy as np from PIL import Image -import utils -import vision_transformer as vits +import dino.utils +import dino.vision_transformer as vits FOURCC = { diff --git a/vision_transformer.py b/dino/vision_transformer.py similarity index 99% rename from vision_transformer.py rename to dino/vision_transformer.py index f69a7ad05..4ad800469 100644 --- a/vision_transformer.py +++ b/dino/vision_transformer.py @@ -21,7 +21,7 @@ import torch import torch.nn as nn -from utils import trunc_normal_ +from dino.utils import trunc_normal_ def drop_path(x, drop_prob: float = 0., training: bool = False): diff --git a/visualize_attention.py b/dino/visualize_attention.py similarity index 99% rename from visualize_attention.py rename to dino/visualize_attention.py index 4288265b9..09891e6de 100644 --- a/visualize_attention.py +++ b/dino/visualize_attention.py @@ -31,8 +31,8 @@ import numpy as np from PIL import Image -import utils -import vision_transformer as vits +import dino.utils +import dino.vision_transformer as vits def apply_mask(image, mask, color, alpha=0.5): diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..e0f54c0ff --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,7 @@ +[build-system] +build-backend = "setuptools.build_meta" +requires = [ "setuptools" ] + +[project] +name = "dino" +version = "2023.04.27"