caizhongang
341 lines
14 KiB
Python
341 lines
14 KiB
Python
import os.path as osp
|
|
import math
|
|
import abc
|
|
from torch.utils.data import DataLoader
|
|
import torch.optim
|
|
import torchvision.transforms as transforms
|
|
from timer import Timer
|
|
from logger import colorlogger
|
|
from torch.nn.parallel.data_parallel import DataParallel
|
|
from config import cfg
|
|
from SMPLer_X import get_model
|
|
from dataset import MultipleDatasets
|
|
# ddp
|
|
import torch.distributed as dist
|
|
from torch.utils.data import DistributedSampler
|
|
import torch.utils.data.distributed
|
|
from utils.distribute_utils import (
|
|
get_rank, is_main_process, time_synchronized, get_group_idx, get_process_groups
|
|
)
|
|
from mmcv.runner import get_dist_info
|
|
|
|
# dynamic dataset import
|
|
for i in range(len(cfg.trainset_3d)):
|
|
exec('from ' + cfg.trainset_3d[i] + ' import ' + cfg.trainset_3d[i])
|
|
for i in range(len(cfg.trainset_2d)):
|
|
exec('from ' + cfg.trainset_2d[i] + ' import ' + cfg.trainset_2d[i])
|
|
for i in range(len(cfg.trainset_humandata)):
|
|
exec('from ' + cfg.trainset_humandata[i] + ' import ' + cfg.trainset_humandata[i])
|
|
exec('from ' + cfg.testset + ' import ' + cfg.testset)
|
|
|
|
|
|
class Base(object):
|
|
__metaclass__ = abc.ABCMeta
|
|
|
|
def __init__(self, log_name='logs.txt'):
|
|
self.cur_epoch = 0
|
|
|
|
# timer
|
|
self.tot_timer = Timer()
|
|
self.gpu_timer = Timer()
|
|
self.read_timer = Timer()
|
|
|
|
# logger
|
|
self.logger = colorlogger(cfg.log_dir, log_name=log_name)
|
|
|
|
@abc.abstractmethod
|
|
def _make_batch_generator(self):
|
|
return
|
|
|
|
@abc.abstractmethod
|
|
def _make_model(self):
|
|
return
|
|
|
|
|
|
class Trainer(Base):
|
|
def __init__(self, distributed=False, gpu_idx=None):
|
|
super(Trainer, self).__init__(log_name='train_logs.txt')
|
|
self.distributed = distributed
|
|
self.gpu_idx = gpu_idx
|
|
|
|
def get_optimizer(self, model):
|
|
normal_param = []
|
|
special_param = []
|
|
for module in model.module.special_trainable_modules:
|
|
special_param += list(module.parameters())
|
|
# print(module)
|
|
for module in model.module.trainable_modules:
|
|
normal_param += list(module.parameters())
|
|
# self.logger.info(f"N-{self.gpu_idx}, {normal_param}")
|
|
# self.logger.info("S", special_param)
|
|
optim_params = [
|
|
{ # add normal params first
|
|
'params': normal_param,
|
|
'lr': cfg.lr
|
|
},
|
|
{
|
|
'params': special_param,
|
|
'lr': cfg.lr * cfg.lr_mult
|
|
},
|
|
]
|
|
optimizer = torch.optim.Adam(optim_params, lr=cfg.lr)
|
|
return optimizer
|
|
|
|
def save_model(self, state, epoch):
|
|
file_path = osp.join(cfg.model_dir, 'snapshot_{}.pth.tar'.format(str(epoch)))
|
|
|
|
# do not save smplx layer weights
|
|
dump_key = []
|
|
for k in state['network'].keys():
|
|
if 'smplx_layer' in k:
|
|
dump_key.append(k)
|
|
for k in dump_key:
|
|
state['network'].pop(k, None)
|
|
|
|
torch.save(state, file_path)
|
|
self.logger.info("Write snapshot into {}".format(file_path))
|
|
|
|
def load_model(self, model, optimizer):
|
|
if cfg.pretrained_model_path is not None:
|
|
ckpt_path = cfg.pretrained_model_path
|
|
ckpt = torch.load(ckpt_path, map_location=torch.device('cpu')) # solve CUDA OOM error in DDP
|
|
model.load_state_dict(ckpt['network'], strict=False)
|
|
self.logger.info('Load checkpoint from {}'.format(ckpt_path))
|
|
if not hasattr(cfg, 'start_over') or cfg.start_over:
|
|
start_epoch = 0
|
|
else:
|
|
optimizer.load_state_dict(ckpt['optimizer'])
|
|
start_epoch = ckpt['epoch'] + 1
|
|
self.logger.info(f'Load optimizer, start from{start_epoch}')
|
|
else:
|
|
start_epoch = 0
|
|
|
|
return start_epoch, model, optimizer
|
|
|
|
def get_lr(self):
|
|
for g in self.optimizer.param_groups:
|
|
cur_lr = g['lr']
|
|
return cur_lr
|
|
|
|
def _make_batch_generator(self):
|
|
# data load and construct batch generator
|
|
self.logger_info("Creating dataset...")
|
|
trainset3d_loader = []
|
|
for i in range(len(cfg.trainset_3d)):
|
|
trainset3d_loader.append(eval(cfg.trainset_3d[i])(transforms.ToTensor(), "train"))
|
|
trainset2d_loader = []
|
|
for i in range(len(cfg.trainset_2d)):
|
|
trainset2d_loader.append(eval(cfg.trainset_2d[i])(transforms.ToTensor(), "train"))
|
|
trainset_humandata_loader = []
|
|
for i in range(len(cfg.trainset_humandata)):
|
|
trainset_humandata_loader.append(eval(cfg.trainset_humandata[i])(transforms.ToTensor(), "train"))
|
|
|
|
data_strategy = getattr(cfg, 'data_strategy', None)
|
|
if data_strategy == 'concat':
|
|
print("Using [concat] strategy...")
|
|
trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader,
|
|
make_same_len=False, verbose=True)
|
|
elif data_strategy == 'balance':
|
|
total_len = getattr(cfg, 'total_data_len', 'auto')
|
|
print(f"Using [balance] strategy with total_data_len : {total_len}...")
|
|
trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader,
|
|
make_same_len=True, total_len=total_len, verbose=True)
|
|
else:
|
|
# original strategy implementation
|
|
valid_loader_num = 0
|
|
if len(trainset3d_loader) > 0:
|
|
trainset3d_loader = [MultipleDatasets(trainset3d_loader, make_same_len=False)]
|
|
valid_loader_num += 1
|
|
else:
|
|
trainset3d_loader = []
|
|
if len(trainset2d_loader) > 0:
|
|
trainset2d_loader = [MultipleDatasets(trainset2d_loader, make_same_len=False)]
|
|
valid_loader_num += 1
|
|
else:
|
|
trainset2d_loader = []
|
|
if len(trainset_humandata_loader) > 0:
|
|
trainset_humandata_loader = [MultipleDatasets(trainset_humandata_loader, make_same_len=False)]
|
|
valid_loader_num += 1
|
|
|
|
if valid_loader_num > 1:
|
|
trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, make_same_len=True)
|
|
else:
|
|
trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, make_same_len=False)
|
|
|
|
self.itr_per_epoch = math.ceil(len(trainset_loader) / cfg.num_gpus / cfg.train_batch_size)
|
|
|
|
if self.distributed:
|
|
self.logger_info(f"Total data length {len(trainset_loader)}.")
|
|
rank, world_size = get_dist_info()
|
|
self.logger_info("Using distributed data sampler.")
|
|
|
|
sampler_train = DistributedSampler(trainset_loader, world_size, rank, shuffle=True)
|
|
self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.train_batch_size,
|
|
shuffle=False, num_workers=cfg.num_thread, sampler=sampler_train,
|
|
pin_memory=True, persistent_workers=True if cfg.num_thread > 0 else False, drop_last=True)
|
|
else:
|
|
self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.num_gpus * cfg.train_batch_size,
|
|
shuffle=True, num_workers=cfg.num_thread,
|
|
pin_memory=True, drop_last=True)
|
|
|
|
def _make_model(self):
|
|
# prepare network
|
|
self.logger_info("Creating graph and optimizer...")
|
|
model = get_model('train')
|
|
|
|
if getattr(cfg, 'fine_tune', None) == 'backbone':
|
|
print("Fine-tuning [backbone]...")
|
|
for module in model.head:
|
|
for param in module.parameters():
|
|
param.requires_grad = False
|
|
for module in model.neck:
|
|
for param in module.parameters():
|
|
param.requires_grad = False
|
|
|
|
elif getattr(cfg, 'fine_tune', None) == 'neck_and_head':
|
|
print("Fine-tuning [neck and head]...")
|
|
for param in model.encoder.parameters():
|
|
param.requires_grad = False
|
|
|
|
elif getattr(cfg, 'fine_tune', None) == 'head':
|
|
print("Fine-tuning [head]...")
|
|
for param in model.encoder.parameters():
|
|
param.requires_grad = False
|
|
for module in model.neck:
|
|
for param in module.parameters():
|
|
param.requires_grad = False
|
|
|
|
|
|
# ddp
|
|
if self.distributed:
|
|
self.logger_info("Using distributed data parallel.")
|
|
model.cuda()
|
|
if hasattr(cfg, 'syncbn') and cfg.syncbn:
|
|
self.logger_info("Using sync batch norm layers.")
|
|
|
|
process_groups = get_process_groups()
|
|
process_group = process_groups[get_group_idx()]
|
|
syncbn_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model, process_group)
|
|
model = torch.nn.parallel.DistributedDataParallel(
|
|
syncbn_model, device_ids=[self.gpu_idx],
|
|
find_unused_parameters=True)
|
|
else:
|
|
model = torch.nn.parallel.DistributedDataParallel(
|
|
model, device_ids=[self.gpu_idx],
|
|
find_unused_parameters=True)
|
|
else:
|
|
# dp
|
|
model = DataParallel(model).cuda()
|
|
|
|
optimizer = self.get_optimizer(model)
|
|
|
|
if hasattr(cfg, "scheduler"):
|
|
if cfg.scheduler == 'cos':
|
|
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cfg.end_epoch * self.itr_per_epoch,
|
|
eta_min=1e-6)
|
|
elif cfg.scheduler == 'step':
|
|
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, cfg.step_size, gamma=cfg.gamma,
|
|
last_epoch=- 1, verbose=False)
|
|
|
|
else:
|
|
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cfg.end_epoch * self.itr_per_epoch,
|
|
eta_min=getattr(cfg,'min_lr',1e-6))
|
|
if cfg.continue_train:
|
|
if self.distributed:
|
|
start_epoch, model, optimizer = self.load_model(model, optimizer)
|
|
else:
|
|
start_epoch, model, optimizer = self.load_model(model, optimizer)
|
|
else:
|
|
start_epoch = 0
|
|
model.train()
|
|
|
|
self.scheduler = scheduler
|
|
self.start_epoch = start_epoch
|
|
self.model = model
|
|
self.optimizer = optimizer
|
|
|
|
def logger_info(self, info):
|
|
if self.distributed:
|
|
if is_main_process():
|
|
self.logger.info(info)
|
|
else:
|
|
self.logger.info(info)
|
|
|
|
|
|
class Tester(Base):
|
|
def __init__(self, test_epoch=None):
|
|
if test_epoch is not None:
|
|
self.test_epoch = int(test_epoch)
|
|
super(Tester, self).__init__(log_name='test_logs.txt')
|
|
|
|
def _make_batch_generator(self):
|
|
# data load and construct batch generator
|
|
self.logger.info("Creating dataset...")
|
|
testset_loader = eval(cfg.testset)(transforms.ToTensor(), "test")
|
|
batch_generator = DataLoader(dataset=testset_loader, batch_size=cfg.num_gpus * cfg.test_batch_size,
|
|
shuffle=False, num_workers=cfg.num_thread, pin_memory=True)
|
|
|
|
self.testset = testset_loader
|
|
self.batch_generator = batch_generator
|
|
|
|
def _make_model(self):
|
|
self.logger.info('Load checkpoint from {}'.format(cfg.pretrained_model_path))
|
|
|
|
# prepare network
|
|
self.logger.info("Creating graph...")
|
|
model = get_model('test')
|
|
model = DataParallel(model).cuda()
|
|
if not getattr(cfg, 'random_init', False):
|
|
ckpt = torch.load(cfg.pretrained_model_path, map_location=torch.device('cpu'))
|
|
|
|
from collections import OrderedDict
|
|
new_state_dict = OrderedDict()
|
|
for k, v in ckpt['network'].items():
|
|
if 'module' not in k:
|
|
k = 'module.' + k
|
|
k = k.replace('backbone', 'encoder').replace('body_rotation_net', 'body_regressor').replace(
|
|
'hand_rotation_net', 'hand_regressor')
|
|
new_state_dict[k] = v
|
|
self.logger.warning("Attention: Strict=False is set for checkpoint loading. Please check manually.")
|
|
model.load_state_dict(new_state_dict, strict=False)
|
|
model.eval()
|
|
else:
|
|
print('Random init!!!!!!!')
|
|
|
|
self.model = model
|
|
|
|
def _evaluate(self, outs, cur_sample_idx):
|
|
eval_result = self.testset.evaluate(outs, cur_sample_idx)
|
|
return eval_result
|
|
|
|
def _print_eval_result(self, eval_result):
|
|
self.testset.print_eval_result(eval_result)
|
|
|
|
class Demoer(Base):
|
|
def __init__(self, test_epoch=None):
|
|
if test_epoch is not None:
|
|
self.test_epoch = int(test_epoch)
|
|
super(Demoer, self).__init__(log_name='test_logs.txt')
|
|
|
|
def _make_model(self):
|
|
self.logger.info('Load checkpoint from {}'.format(cfg.pretrained_model_path))
|
|
|
|
# prepare network
|
|
self.logger.info("Creating graph...")
|
|
model = get_model('test')
|
|
model = DataParallel(model).cuda()
|
|
ckpt = torch.load(cfg.pretrained_model_path, map_location=torch.device('cpu'))
|
|
|
|
from collections import OrderedDict
|
|
new_state_dict = OrderedDict()
|
|
for k, v in ckpt['network'].items():
|
|
if 'module' not in k:
|
|
k = 'module.' + k
|
|
k = k.replace('module.backbone', 'module.encoder').replace('body_rotation_net', 'body_regressor').replace(
|
|
'hand_rotation_net', 'hand_regressor')
|
|
new_state_dict[k] = v
|
|
model.load_state_dict(new_state_dict, strict=False)
|
|
model.eval()
|
|
|
|
self.model = model
|