# 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. # -------------------------------------------------------- # References: # 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 time from pathlib import Path import torch import torch.backends.cudnn as cudnn from torch.utils.tensorboard import SummaryWriter import torchvision.transforms as transforms import torchvision.datasets as datasets import resource import timm import timm.optim.optim_factory as optim_factory import util.misc as misc from util.misc import NativeScalerWithGradNormCount as NativeScaler import models_base_ofa_mae as models_mae #import models_large_ofa_mae as models_mae from engine_pretrain import train_one_epoch from datasets.ofall_dataset import Sentinel1Dataset, Sentinel2Dataset, NAIPDataset, HyperDataset, GaufenDataset, MyBatchSampler, MyDistributedBatchSampler def get_args_parser(): parser = argparse.ArgumentParser('MAE pre-training', add_help=False) 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=400, type=int) parser.add_argument('--accum_iter', default=1, type=int, help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)') # Model parameters parser.add_argument('--model', default='mae_vit_large_patch16', type=str, metavar='MODEL', help='Name of model to train') parser.add_argument('--input_size', default=224, type=int, help='images input size') parser.add_argument('--mask_ratio', default=0.75, type=float, help='Masking ratio (percentage of removed patches).') parser.add_argument('--norm_pix_loss', action='store_true', help='Use (per-patch) normalized pixels as targets for computing loss') parser.set_defaults(norm_pix_loss=False) # Optimizer parameters 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=1e-3, 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') # Dataset parameters parser.add_argument('--data_path_s1', type=str, help='dataset path s1') parser.add_argument('--data_path_s2', type=str, help='dataset path s2') parser.add_argument('--data_path_naip', type=str, help='dataset path naip') parser.add_argument('--data_path_hyper', type=str, help='dataset path hyper') parser.add_argument('--data_path_gaufen', type=str, help='dataset path gaufen') parser.add_argument('--output_dir', default='./output_dir', help='path where to save, empty for no saving') parser.add_argument('--log_dir', default='./output_dir', 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('--resume', default='', help='resume from checkpoint') parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch') parser.add_argument('--num_workers', default=10, type=int) parser.add_argument('--pin_mem', action='store_true', help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.') parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem') parser.set_defaults(pin_mem=True) # 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') parser.add_argument('--in_chans', nargs='*', default=[2,9,3,70,4], type=int) parser.add_argument("--local-rank", type=int, default=0) parser.add_argument("--split", type=str, default='train_10k') parser.add_argument("--wave_list", type=list, default=None) return parser def main(args): misc.init_distributed_mode(args) print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__)))) print("{}".format(args).replace(', ', ',\n')) device = torch.device(args.device) # fix the seed for reproducibility seed = args.seed + misc.get_rank() torch.manual_seed(seed) np.random.seed(seed) cudnn.benchmark = True if True: print(f'Loading data with split: {args.split}') dataset_s1 = Sentinel1Dataset(root_dir=args.data_path_s1, split=args.split) dataset_s2 = Sentinel2Dataset(root_dir=args.data_path_s2, split=args.split) dataset_naip = NAIPDataset(root_dir=args.data_path_naip, split=args.split) dataset_hyper = HyperDataset(root_dir=args.data_path_hyper, split='train_10k') dataset_gaufen = GaufenDataset(root_dir=args.data_path_gaufen, split=args.split) dataset_train = torch.utils.data.ConcatDataset([dataset_s1,dataset_s2,dataset_naip,dataset_hyper,dataset_gaufen]) else: raise NotImplementedError #pdb.set_trace() if True: # args.distributed: num_tasks = misc.get_world_size() global_rank = misc.get_rank() batch_sampler_train = MyDistributedBatchSampler( dataset_train, dataset_s1, dataset_s2, dataset_naip, dataset_hyper, dataset_gaufen, num_replicas=num_tasks, rank=global_rank, shuffle=True, batch_size = args.batch_size, drop_last=True ) print("Sampler_train = %s" % str(batch_sampler_train)) else: batch_sampler_train = MyBatchSampler( dataset_train, dataset_s1, dataset_s2, dataset_naip, dataset_hyper, dataset_gaufen, batch_size=args.batch_size, shuffle=True, drop_last=True ) if global_rank == 0 and args.log_dir is not None: os.makedirs(args.log_dir, exist_ok=True) log_writer = SummaryWriter(log_dir=args.log_dir) else: log_writer = None data_loader_train = torch.utils.data.DataLoader( dataset_train, batch_sampler=batch_sampler_train, num_workers=args.num_workers, pin_memory=args.pin_mem, ) print(len(dataset_train),len(data_loader_train)) #pdb.set_trace() # define the model model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss, in_chans=args.in_chans) 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) model_without_ddp = model.module # following timm: set wd as 0 for bias and norm layers param_groups = optim_factory.param_groups_weight_decay(model_without_ddp, args.weight_decay) # timm 1.x optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95)) loss_scaler = NativeScaler() misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler) print(f"Start training for {args.epochs} epochs") start_time = time.time() for epoch in range(args.start_epoch, args.epochs): if args.distributed: data_loader_train.batch_sampler.set_epoch(epoch) train_stats = train_one_epoch( model, data_loader_train, optimizer, device, epoch, loss_scaler, log_writer=log_writer, args=args ) if args.output_dir and (epoch % 10 == 0 or epoch + 1 == args.epochs): misc.save_model( args=args, model=model, 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)) if __name__ == '__main__': args = get_args_parser() args = args.parse_args() if args.output_dir: Path(args.output_dir).mkdir(parents=True, exist_ok=True) main(args)