wehub-resource-sync
230 lines
8.6 KiB
Python
230 lines
8.6 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import glob
|
|
import argparse
|
|
from easydict import EasyDict as edict
|
|
|
|
import torch
|
|
import torch.multiprocessing as mp
|
|
import numpy as np
|
|
import random
|
|
|
|
try:
|
|
import wandb
|
|
WANDB_AVAILABLE = True
|
|
except ImportError:
|
|
WANDB_AVAILABLE = False
|
|
print("Warning: wandb not installed. Install with 'pip install wandb' to enable wandb logging.")
|
|
|
|
from pixal3d import models, datasets, trainers
|
|
from pixal3d.utils.dist_utils import setup_dist
|
|
|
|
|
|
def find_ckpt(cfg):
|
|
# Load checkpoint
|
|
cfg['load_ckpt'] = None
|
|
if cfg.load_dir != '':
|
|
if cfg.ckpt == 'latest':
|
|
files = glob.glob(os.path.join(cfg.load_dir, 'ckpts', 'misc_*.pt'))
|
|
if len(files) != 0:
|
|
cfg.load_ckpt = max([
|
|
int(os.path.basename(f).split('step')[-1].split('.')[0])
|
|
for f in files
|
|
])
|
|
elif cfg.ckpt == 'none':
|
|
cfg.load_ckpt = None
|
|
else:
|
|
cfg.load_ckpt = int(cfg.ckpt)
|
|
return cfg
|
|
|
|
|
|
def setup_rng(rank):
|
|
torch.manual_seed(rank)
|
|
torch.cuda.manual_seed_all(rank)
|
|
np.random.seed(rank)
|
|
random.seed(rank)
|
|
|
|
|
|
def get_model_summary(model):
|
|
num_params = 0
|
|
num_trainable_params = 0
|
|
for name, param in model.named_parameters():
|
|
num_params += param.numel()
|
|
if param.requires_grad:
|
|
num_trainable_params += param.numel()
|
|
model_summary = f'Number of parameters: {num_params:,}\n'
|
|
model_summary += f'Number of trainable parameters: {num_trainable_params:,}\n'
|
|
return model_summary
|
|
|
|
|
|
def main(local_rank, cfg):
|
|
# Set up distributed training
|
|
rank = cfg.node_rank * cfg.num_gpus + local_rank
|
|
world_size = cfg.num_nodes * cfg.num_gpus
|
|
if world_size > 1:
|
|
setup_dist(rank, local_rank, world_size, cfg.master_addr, cfg.master_port)
|
|
|
|
# Multi-GPU training verification
|
|
print(f'[Rank {rank}/{world_size}] Process started on GPU {local_rank} (cuda:{torch.cuda.current_device()})')
|
|
if rank == 0:
|
|
print(f'\n{"="*60}')
|
|
print(f'Multi-GPU Training Verification:')
|
|
print(f' - Total GPUs (world_size): {world_size}')
|
|
print(f' - num_gpus per node: {cfg.num_gpus}')
|
|
print(f' - num_nodes: {cfg.num_nodes}')
|
|
print(f'{"="*60}\n')
|
|
|
|
# Initialize wandb (only on rank 0)
|
|
wandb_run = None
|
|
if rank == 0 and cfg.use_wandb and WANDB_AVAILABLE:
|
|
# Use WANDB_DIR env var (local SSD) to avoid S3 FUSE rename/append issues
|
|
wandb_dir = os.environ.get('WANDB_DIR', cfg.output_dir)
|
|
os.makedirs(wandb_dir, exist_ok=True)
|
|
wandb_run = wandb.init(
|
|
project=cfg.wandb_project,
|
|
name=cfg.wandb_name if cfg.wandb_name else os.path.basename(cfg.output_dir),
|
|
config=cfg.__dict__,
|
|
dir=wandb_dir,
|
|
resume="allow",
|
|
id=cfg.wandb_id if cfg.wandb_id else None,
|
|
)
|
|
print(f'Wandb initialized: {wandb_run.url}')
|
|
|
|
# Upload config JSON file as wandb artifact
|
|
config_file = cfg.get('config', None)
|
|
if config_file and os.path.isfile(config_file):
|
|
config_artifact = wandb.Artifact(
|
|
name=f"config-{wandb_run.id}",
|
|
type="config",
|
|
description=f"Training config for {wandb_run.name}",
|
|
)
|
|
config_artifact.add_file(config_file, name=os.path.basename(config_file))
|
|
# Also save the resolved full config (with command-line args merged)
|
|
resolved_config_path = os.path.join(cfg.output_dir, 'config_resolved.json')
|
|
with open(resolved_config_path, 'w') as f:
|
|
json.dump(cfg.__dict__, f, indent=4, default=str)
|
|
config_artifact.add_file(resolved_config_path, name='config_resolved.json')
|
|
wandb_run.log_artifact(config_artifact)
|
|
print(f'Uploaded config artifact to wandb: {config_file}')
|
|
|
|
# Seed rngs
|
|
setup_rng(rank)
|
|
|
|
# Load data
|
|
dataset_kwargs = dict(cfg.dataset.args)
|
|
dataset = getattr(datasets, cfg.dataset.name)(cfg.data_dir, **dataset_kwargs)
|
|
|
|
# Print dataset info (only on rank 0)
|
|
if rank == 0:
|
|
print(f'\nDataset: {cfg.dataset.name}, Number of samples: {len(dataset):,}\n')
|
|
|
|
# Build model
|
|
model_dict = {
|
|
name: getattr(models, model.name)(**model.args).cuda()
|
|
for name, model in cfg.models.items()
|
|
}
|
|
|
|
# Model summary
|
|
if rank == 0:
|
|
for name, backbone in model_dict.items():
|
|
model_summary = get_model_summary(backbone)
|
|
print(f'\n\nBackbone: {name}\n' + model_summary)
|
|
with open(os.path.join(cfg.output_dir, f'{name}_model_summary.txt'), 'w') as fp:
|
|
print(model_summary, file=fp)
|
|
|
|
# Build trainer
|
|
trainer = getattr(trainers, cfg.trainer.name)(
|
|
model_dict, dataset,
|
|
**cfg.trainer.args,
|
|
output_dir=cfg.output_dir,
|
|
load_dir=cfg.load_dir,
|
|
step=cfg.load_ckpt,
|
|
wandb_run=wandb_run, # Pass wandb run to trainer
|
|
)
|
|
|
|
# Train
|
|
if not cfg.tryrun:
|
|
if cfg.profile:
|
|
trainer.profile()
|
|
else:
|
|
trainer.run()
|
|
|
|
# Close wandb
|
|
if wandb_run is not None:
|
|
wandb_run.finish()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# Arguments and config
|
|
parser = argparse.ArgumentParser()
|
|
## config
|
|
parser.add_argument('--config', type=str, required=True, help='Experiment config file')
|
|
## io and resume
|
|
parser.add_argument('--output_dir', type=str, required=True, help='Output directory')
|
|
parser.add_argument('--load_dir', type=str, default='', help='Load directory, default to output_dir')
|
|
parser.add_argument('--ckpt', type=str, default='latest', help='Checkpoint step to resume training, default to latest')
|
|
parser.add_argument('--data_dir', type=str, default='./data/', help='Data directory')
|
|
parser.add_argument('--auto_retry', type=int, default=3, help='Number of retries on error')
|
|
## dubug
|
|
parser.add_argument('--tryrun', action='store_true', help='Try run without training')
|
|
parser.add_argument('--profile', action='store_true', help='Profile training')
|
|
## multi-node and multi-gpu
|
|
parser.add_argument('--num_nodes', type=int, default=1, help='Number of nodes')
|
|
parser.add_argument('--node_rank', type=int, default=0, help='Node rank')
|
|
parser.add_argument('--num_gpus', type=int, default=-1, help='Number of GPUs per node, default to all')
|
|
parser.add_argument('--master_addr', type=str, default='localhost', help='Master address for distributed training')
|
|
parser.add_argument('--master_port', type=str, default='12666', help='Port for distributed training')
|
|
## wandb
|
|
parser.add_argument('--use_wandb', action='store_true', help='Enable wandb logging')
|
|
parser.add_argument('--wandb_project', type=str, default='pixal3d-training', help='Wandb project name')
|
|
parser.add_argument('--wandb_name', type=str, default='', help='Wandb run name, default to output_dir basename')
|
|
parser.add_argument('--wandb_id', type=str, default='', help='Wandb run id for resuming')
|
|
opt = parser.parse_args()
|
|
opt.load_dir = opt.load_dir if opt.load_dir != '' else opt.output_dir
|
|
opt.num_gpus = torch.cuda.device_count() if opt.num_gpus == -1 else opt.num_gpus
|
|
## Load config
|
|
config = json.load(open(opt.config, 'r'))
|
|
## Combine arguments and config
|
|
cfg = edict()
|
|
cfg.update(opt.__dict__)
|
|
cfg.update(config)
|
|
print('\n\nConfig:')
|
|
print('=' * 80)
|
|
print(json.dumps(cfg.__dict__, indent=4))
|
|
|
|
# Prepare output directory
|
|
if cfg.node_rank == 0:
|
|
os.makedirs(cfg.output_dir, exist_ok=True)
|
|
## Save command and config
|
|
with open(os.path.join(cfg.output_dir, 'command.txt'), 'w') as fp:
|
|
print(' '.join(['python'] + sys.argv), file=fp)
|
|
with open(os.path.join(cfg.output_dir, 'config.json'), 'w') as fp:
|
|
json.dump(config, fp, indent=4)
|
|
|
|
# Run
|
|
if cfg.auto_retry == 0:
|
|
cfg = find_ckpt(cfg)
|
|
if cfg.num_gpus > 1:
|
|
mp.spawn(main, args=(cfg,), nprocs=cfg.num_gpus, join=True)
|
|
else:
|
|
main(0, cfg)
|
|
else:
|
|
for rty in range(cfg.auto_retry):
|
|
try:
|
|
cfg = find_ckpt(cfg)
|
|
if cfg.num_gpus > 1:
|
|
mp.spawn(main, args=(cfg,), nprocs=cfg.num_gpus, join=True)
|
|
else:
|
|
main(0, cfg)
|
|
break
|
|
except Exception as e:
|
|
import traceback
|
|
print(f'\n{"="*60}')
|
|
print(f'Error: {e}')
|
|
print(f'{"="*60}')
|
|
print('Full traceback:')
|
|
traceback.print_exc()
|
|
print(f'{"="*60}')
|
|
print(f'Retrying ({rty + 1}/{cfg.auto_retry})...')
|
|
|