StableRecon / croco /pretrain.py
Stable-X's picture
Upload folder using huggingface_hub
e4bf056 verified
raw
history blame
11.5 kB
# Copyright (C) 2022-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# Pre-training CroCo
# --------------------------------------------------------
# References:
# MAE: https://github.com/facebookresearch/mae
# DeiT: https://github.com/facebookresearch/deit
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------
import argparse
import datetime
import json
import numpy as np
import os
import sys
import time
import math
from pathlib import Path
from typing import Iterable
import torch
import torch.distributed as dist
import torch.backends.cudnn as cudnn
from torch.utils.tensorboard import SummaryWriter
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import utils.misc as misc
from utils.misc import NativeScalerWithGradNormCount as NativeScaler
from models.croco import CroCoNet
from models.criterion import MaskedMSE
from datasets.pairs_dataset import PairsDataset
def get_args_parser():
parser = argparse.ArgumentParser('CroCo pre-training', add_help=False)
# model and criterion
parser.add_argument('--model', default='CroCoNet()', type=str, help="string containing the model to build")
parser.add_argument('--norm_pix_loss', default=1, choices=[0,1], help="apply per-patch mean/std normalization before applying the loss")
# dataset
parser.add_argument('--dataset', default='habitat_release', type=str, help="training set")
parser.add_argument('--transforms', default='crop224+acolor', type=str, help="transforms to apply") # in the paper, we also use some homography and rotation, but find later that they were not useful or even harmful
# training
parser.add_argument('--seed', default=0, type=int, help="Random seed")
parser.add_argument('--batch_size', default=64, type=int, help="Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus")
parser.add_argument('--epochs', default=800, type=int, help="Maximum number of epochs for the scheduler")
parser.add_argument('--max_epoch', default=400, type=int, help="Stop training at this epoch")
parser.add_argument('--accum_iter', default=1, type=int, help="Accumulate gradient iterations (for increasing the effective batch size under memory constraints)")
parser.add_argument('--weight_decay', type=float, default=0.05, help="weight decay (default: 0.05)")
parser.add_argument('--lr', type=float, default=None, metavar='LR', help='learning rate (absolute lr)')
parser.add_argument('--blr', type=float, default=1.5e-4, metavar='LR', help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
parser.add_argument('--min_lr', type=float, default=0., metavar='LR', help='lower lr bound for cyclic schedulers that hit 0')
parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N', help='epochs to warmup LR')
parser.add_argument('--amp', type=int, default=1, choices=[0,1], help="Use Automatic Mixed Precision for pretraining")
# others
parser.add_argument('--num_workers', default=8, type=int)
parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes')
parser.add_argument('--local_rank', default=-1, type=int)
parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
parser.add_argument('--save_freq', default=1, type=int, help='frequence (number of epochs) to save checkpoint in checkpoint-last.pth')
parser.add_argument('--keep_freq', default=20, type=int, help='frequence (number of epochs) to save checkpoint in checkpoint-%d.pth')
parser.add_argument('--print_freq', default=20, type=int, help='frequence (number of iterations) to print infos while training')
# paths
parser.add_argument('--output_dir', default='./output/', type=str, help="path where to save the output")
parser.add_argument('--data_dir', default='./data/', type=str, help="path where data are stored")
return parser
def main(args):
misc.init_distributed_mode(args)
global_rank = misc.get_rank()
world_size = misc.get_world_size()
print("output_dir: "+args.output_dir)
if args.output_dir:
Path(args.output_dir).mkdir(parents=True, exist_ok=True)
# auto resume
last_ckpt_fname = os.path.join(args.output_dir, f'checkpoint-last.pth')
args.resume = last_ckpt_fname if os.path.isfile(last_ckpt_fname) else None
print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
print("{}".format(args).replace(', ', ',\n'))
device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
# fix the seed
seed = args.seed + misc.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
cudnn.benchmark = True
## training dataset and loader
print('Building dataset for {:s} with transforms {:s}'.format(args.dataset, args.transforms))
dataset = PairsDataset(args.dataset, trfs=args.transforms, data_dir=args.data_dir)
if world_size>1:
sampler_train = torch.utils.data.DistributedSampler(
dataset, num_replicas=world_size, rank=global_rank, shuffle=True
)
print("Sampler_train = %s" % str(sampler_train))
else:
sampler_train = torch.utils.data.RandomSampler(dataset)
data_loader_train = torch.utils.data.DataLoader(
dataset, sampler=sampler_train,
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True,
)
## model
print('Loading model: {:s}'.format(args.model))
model = eval(args.model)
print('Loading criterion: MaskedMSE(norm_pix_loss={:s})'.format(str(bool(args.norm_pix_loss))))
criterion = MaskedMSE(norm_pix_loss=bool(args.norm_pix_loss))
model.to(device)
model_without_ddp = model
print("Model = %s" % str(model_without_ddp))
eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
if args.lr is None: # only base_lr is specified
args.lr = args.blr * eff_batch_size / 256
print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
print("actual lr: %.2e" % args.lr)
print("accumulate grad iterations: %d" % args.accum_iter)
print("effective batch size: %d" % eff_batch_size)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True, static_graph=True)
model_without_ddp = model.module
param_groups = misc.get_parameter_groups(model_without_ddp, args.weight_decay) # following timm: set wd as 0 for bias and norm layers
optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
print(optimizer)
loss_scaler = NativeScaler()
misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
if global_rank == 0 and args.output_dir is not None:
log_writer = SummaryWriter(log_dir=args.output_dir)
else:
log_writer = None
print(f"Start training until {args.max_epoch} epochs")
start_time = time.time()
for epoch in range(args.start_epoch, args.max_epoch):
if world_size>1:
data_loader_train.sampler.set_epoch(epoch)
train_stats = train_one_epoch(
model, criterion, data_loader_train,
optimizer, device, epoch, loss_scaler,
log_writer=log_writer,
args=args
)
if args.output_dir and epoch % args.save_freq == 0 :
misc.save_model(
args=args, model_without_ddp=model_without_ddp, optimizer=optimizer,
loss_scaler=loss_scaler, epoch=epoch, fname='last')
if args.output_dir and (epoch % args.keep_freq == 0 or epoch + 1 == args.max_epoch) and (epoch>0 or args.max_epoch==1):
misc.save_model(
args=args, model_without_ddp=model_without_ddp, optimizer=optimizer,
loss_scaler=loss_scaler, epoch=epoch)
log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
'epoch': epoch,}
if args.output_dir and misc.is_main_process():
if log_writer is not None:
log_writer.flush()
with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
f.write(json.dumps(log_stats) + "\n")
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
data_loader: Iterable, optimizer: torch.optim.Optimizer,
device: torch.device, epoch: int, loss_scaler,
log_writer=None,
args=None):
model.train(True)
metric_logger = misc.MetricLogger(delimiter=" ")
metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
header = 'Epoch: [{}]'.format(epoch)
accum_iter = args.accum_iter
optimizer.zero_grad()
if log_writer is not None:
print('log_dir: {}'.format(log_writer.log_dir))
for data_iter_step, (image1, image2) in enumerate(metric_logger.log_every(data_loader, args.print_freq, header)):
# we use a per iteration lr scheduler
if data_iter_step % accum_iter == 0:
misc.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
image1 = image1.to(device, non_blocking=True)
image2 = image2.to(device, non_blocking=True)
with torch.cuda.amp.autocast(enabled=bool(args.amp)):
out, mask, target = model(image1, image2)
loss = criterion(out, mask, target)
loss_value = loss.item()
if not math.isfinite(loss_value):
print("Loss is {}, stopping training".format(loss_value))
sys.exit(1)
loss /= accum_iter
loss_scaler(loss, optimizer, parameters=model.parameters(),
update_grad=(data_iter_step + 1) % accum_iter == 0)
if (data_iter_step + 1) % accum_iter == 0:
optimizer.zero_grad()
torch.cuda.synchronize()
metric_logger.update(loss=loss_value)
lr = optimizer.param_groups[0]["lr"]
metric_logger.update(lr=lr)
loss_value_reduce = misc.all_reduce_mean(loss_value)
if log_writer is not None and ((data_iter_step + 1) % (accum_iter*args.print_freq)) == 0:
# x-axis is based on epoch_1000x in the tensorboard, calibrating differences curves when batch size changes
epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
log_writer.add_scalar('train_loss', loss_value_reduce, epoch_1000x)
log_writer.add_scalar('lr', lr, epoch_1000x)
# gather the stats from all processes
metric_logger.synchronize_between_processes()
print("Averaged stats:", metric_logger)
return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
if __name__ == '__main__':
args = get_args_parser()
args = args.parse_args()
main(args)