import argparse import datetime import numpy as np import random import time import torch import json import os from pathlib import Path from optim_factory import create_optimizer from torch.nn.parallel import DistributedDataParallel as DDP from utils import NativeScalerWithGradNormCount as NativeScaler import utils from cwm.data.dataset_utils import build_pretraining_dataset from cwm.model import model_pretrain from engine_for_pretraining import train_one_epoch import wandb import torch.backends.cudnn as cudnn np.random.seed(0) random.seed(0) def get_args(): parser = argparse.ArgumentParser('CWM pre-training script', add_help=False) # training parameters parser.add_argument('--batch_size', default=64, type=int, help='per-GPU batch-size') parser.add_argument('--epochs', default=800, type=int, help='number of training epochs') parser.add_argument('--save_ckpt_freq', default=50, type=int, help='save checkpoint frequency') parser.add_argument('--print_freq', default=1, type=int, help='frequency of printing training stats') parser.add_argument('--accum_iter', default=1, type=int, help='number of steps to accumulate gradients') parser.add_argument('--eval', action='store_true', help='evaluation mode') parser.add_argument('--output_dir', default='', help='path where to save, empty for no saving') parser.add_argument('--log_dir', default=None, help='path where to tensorboard log') parser.add_argument('--device', default='cuda', help='device to use for training / testing') parser.add_argument('--seed', default=0, type=int) parser.add_argument('--val_after', default=50, type=int) parser.add_argument('--resume', default='', help='resume from checkpoint') parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch') parser.add_argument('--use_wandb', action='store_true', help='use wandb for logging') # Model parameters parser.add_argument('--model', default='vitb_8x8patch_3frames', type=str, help='Name of model to train') parser.add_argument('--context_frames', type=int, default=2, help='number of frames model will see densely') parser.add_argument('--target_frames', type=int, default=1, help='number of frames model will see sparsely') parser.add_argument('--temporal_units', type=str, default='ms', help='the units in which time is defined') parser.add_argument('--sampling_rate', type=int, default=150, help='temporal gap between context/target frames') parser.add_argument('--context_target_gap', type=int, nargs='+', default=[150, 150], help='gap between context/target') # Masking and target parameters parser.add_argument('--mask_type', default='rotated_table', type=str, help='masked strategy') parser.add_argument('--mask_ratio', default=0.75, type=float, help='masking ratio') parser.add_argument('--mask_kwargs', default='', type=json.loads, help='extra arguments for masking generator') parser.add_argument('--drop_path', type=float, default=0.0, metavar='PCT', help='Drop path rate (default: 0.1)') # Optimizer parameters parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER', help='Optimizer (default:adamw)') parser.add_argument('--opt_eps', default=1e-8, type=float, metavar='EPSILON', help='Optimizer epsilon') parser.add_argument('--opt_betas', default=None, type=float, nargs='+', metavar='BETA', help='Optimizer Betas') parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='SGD momentum (default: 0.9)') parser.add_argument('--weight_decay', type=float, default=0.05, help='weight decay (default: 0.05)') parser.add_argument('--weight_decay_end', type=float, default=0.05, help='Final value of the weight decay.') parser.add_argument('--lr', type=float, default=1.5e-4, metavar='LR', help='learning rate (default: 1.5e-4)') parser.add_argument('--warmup_lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate') parser.add_argument('--min_lr', type=float, default=0, metavar='LR', help='lower lr bound for cyclic schedulers)') parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N', help='epochs to warmup LR') parser.add_argument('--warmup_steps', type=int, default=-1, metavar='N', help='steps to warmup LR') # Dataset parameters parser.add_argument('--data_path', default='/path/to/list_kinetics-400', type=str, help='dataset path') parser.add_argument('--data_path_list', type=str, nargs='+', default=None, help='[path1, path2, path3, ...]') parser.add_argument('--num_workers', default=10, type=int) # Augmentation parameters parser.add_argument('--augmentation_type', type=str, default='multiscale', choices=['multiscale', 'center', 'none']) parser.add_argument('--augmentation_scales', type=float, nargs='+', default=[1.0, 0.875, 0.75, 0.66]) # distributed training parameters 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_on_itp', action='store_true') parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training') return parser.parse_args() # Assuming 'model' is your PyTorch model def export_model_parameters(model): with open('model_parameters.txt', 'w') as f: for name, param in model.named_parameters(): f.write(f"{name} {param.size()}\n") def main(args): ## Setup distributed training utils.init_distributed_mode(args) cudnn.benchmark = True device = torch.device(args.device) num_tasks = utils.get_world_size() sampler_rank = global_rank = utils.get_rank() world_size = utils.get_world_size() ## Fix the seed for reproducibility seed = args.seed + utils.get_rank() torch.manual_seed(seed) np.random.seed(seed) ## Initialize model model = getattr(model_pretrain, args.model)() args.input_size = int(model.encoder.patch_embed.img_size[0]) args.tubelet_size = model.patch_size[0] args.mask_input_size = ( (args.context_frames + args.target_frames) // args.tubelet_size, args.input_size // model.patch_size[-2], args.input_size // model.patch_size[-1], ) ## Prepare datasets dataset_train = build_pretraining_dataset(args) sampler_train = torch.utils.data.DistributedSampler( dataset_train, num_replicas=num_tasks, rank=sampler_rank, shuffle=True, drop_last=True ) data_loader_train = torch.utils.data.DataLoader( dataset_train, sampler=sampler_train, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True, drop_last=True, worker_init_fn=utils.seed_worker, ) num_steps_per_epoch = len(dataset_train) // args.batch_size // num_tasks n_params, n_params_str = utils.get_model_num_parameters(model) total_batch_size = args.batch_size * world_size * args.accum_iter ## LR and warmup export_model_parameters(model) model = DDP(model.to(device), device_ids=[args.gpu], find_unused_parameters=False) ## Optimizer, loss scaler optimizer = create_optimizer(args, model.module) loss_scaler = NativeScaler() ## LR scheduler, WD scheduler args.lr = args.lr * total_batch_size / 256 args.min_lr = args.min_lr * total_batch_size / 256 args.warmup_lr = args.warmup_lr * total_batch_size / 256 lr_schedule_values = utils.cosine_scheduler( args.lr, args.min_lr, args.epochs, num_steps_per_epoch, warmup_epochs=args.warmup_epochs, warmup_steps=args.warmup_steps, ) wd_schedule_values = utils.cosine_scheduler( args.weight_decay, args.weight_decay_end, args.epochs, num_steps_per_epoch ) ## Resume from checkpoint, if any utils.auto_load_model(args=args, model=model, optimizer=optimizer, loss_scaler=loss_scaler) ## Print training arguments print("world size: %d" % args.world_size) print("model: %s" % args.model) print("image size: %s" % str(args.input_size)) print("patch size: %s" % str(model.module.encoder.patch_embed.patch_size[-2:])) print("context frames: %s" % str(args.context_frames)) print("target frames: %s" % str(args.target_frames)) print("per-device batch size: %d" % total_batch_size) print("total batch size: %d" % total_batch_size) print("grad accumulation: %d" % args.accum_iter) print("dataset length: %d" % len(dataset_train)) print("steps per epoch: %d" % num_steps_per_epoch) print("num parameters: %s" % n_params_str) print("lr: %.8f" % args.lr) ## Setup logging if args.use_wandb and utils.is_main_process(): wandb.init(project="cwm", name=args.output_dir.split('/')[-1], config=args) print(f'start training at epoch {args.start_epoch} for {args.epochs} epochs') start_time = time.time() for epoch in range(args.start_epoch, args.epochs): if args.distributed: data_loader_train.sampler.set_epoch(epoch) # Run one epoch train_stats = train_one_epoch( model, data_loader_train, optimizer, device, epoch, loss_scaler, start_steps=epoch * num_steps_per_epoch, lr_schedule_values=lr_schedule_values, wd_schedule_values=wd_schedule_values, args=args, global_rank=global_rank, ) # Save checkpoint if args.output_dir and ((epoch + 1) % args.save_ckpt_freq == 0 or epoch + 1 == args.epochs): utils.save_model(args=args, model=model, optimizer=optimizer, loss_scaler=loss_scaler, epoch=epoch) # Logging start_time = time.time() do_write = (global_rank == 0) if args.use_xla else utils.is_main_process() if args.output_dir and do_write: log_stats = { **{f'train/{k}': v for k, v in train_stats.items()}, 'epoch': epoch, 'params': n_params, 'epoch_time': time.time() - start_time } with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f: f.write(json.dumps(log_stats) + "\n") if args.use_wandb: wandb.log(log_stats) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) if __name__ == '__main__': opts = get_args() if opts.output_dir: Path(opts.output_dir).mkdir(parents=True, exist_ok=True) main(opts)