File size: 5,834 Bytes
31ca7a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 |
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
from trellis import models, datasets, trainers
from trellis.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):
model_summary = 'Parameters:\n'
model_summary += '=' * 128 + '\n'
model_summary += f'{"Name":<{72}}{"Shape":<{32}}{"Type":<{16}}{"Grad"}\n'
num_params = 0
num_trainable_params = 0
for name, param in model.named_parameters():
model_summary += f'{name:<{72}}{str(param.shape):<{32}}{str(param.dtype):<{16}}{param.requires_grad}\n'
num_params += param.numel()
if param.requires_grad:
num_trainable_params += param.numel()
model_summary += '\n'
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)
# Seed rngs
setup_rng(rank)
# Load data
dataset = getattr(datasets, cfg.dataset.name)(cfg.data_dir, **cfg.dataset.args)
# 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)
# Train
if not cfg.tryrun:
if cfg.profile:
trainer.profile()
else:
trainer.run()
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='12345', help='Port for distributed training')
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:
print(f'Error: {e}')
print(f'Retrying ({rty + 1}/{cfg.auto_retry})...')
|