Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) Open-MMLab. All rights reserved.
|
||||
from .checkpoint import save_checkpoint
|
||||
from .apex_iter_based_runner import IterBasedRunnerAmp
|
||||
|
||||
|
||||
__all__ = [
|
||||
'save_checkpoint', 'IterBasedRunnerAmp',
|
||||
]
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
# Copyright (c) Open-MMLab. All rights reserved.
|
||||
import os.path as osp
|
||||
import platform
|
||||
import shutil
|
||||
|
||||
import torch
|
||||
from torch.optim import Optimizer
|
||||
|
||||
import mmcv
|
||||
from mmcv.runner import RUNNERS, IterBasedRunner
|
||||
from .checkpoint import save_checkpoint
|
||||
|
||||
try:
|
||||
import apex
|
||||
except:
|
||||
print('apex is not installed')
|
||||
|
||||
|
||||
@RUNNERS.register_module()
|
||||
class IterBasedRunnerAmp(IterBasedRunner):
|
||||
"""Iteration-based Runner with AMP support.
|
||||
|
||||
This runner train models iteration by iteration.
|
||||
"""
|
||||
|
||||
def save_checkpoint(self,
|
||||
out_dir,
|
||||
filename_tmpl='iter_{}.pth',
|
||||
meta=None,
|
||||
save_optimizer=True,
|
||||
create_symlink=False):
|
||||
"""Save checkpoint to file.
|
||||
|
||||
Args:
|
||||
out_dir (str): Directory to save checkpoint files.
|
||||
filename_tmpl (str, optional): Checkpoint file template.
|
||||
Defaults to 'iter_{}.pth'.
|
||||
meta (dict, optional): Metadata to be saved in checkpoint.
|
||||
Defaults to None.
|
||||
save_optimizer (bool, optional): Whether save optimizer.
|
||||
Defaults to True.
|
||||
create_symlink (bool, optional): Whether create symlink to the
|
||||
latest checkpoint file. Defaults to True.
|
||||
"""
|
||||
if meta is None:
|
||||
meta = dict(iter=self.iter + 1, epoch=self.epoch + 1)
|
||||
elif isinstance(meta, dict):
|
||||
meta.update(iter=self.iter + 1, epoch=self.epoch + 1)
|
||||
else:
|
||||
raise TypeError(
|
||||
f'meta should be a dict or None, but got {type(meta)}')
|
||||
if self.meta is not None:
|
||||
meta.update(self.meta)
|
||||
|
||||
filename = filename_tmpl.format(self.iter + 1)
|
||||
filepath = osp.join(out_dir, filename)
|
||||
optimizer = self.optimizer if save_optimizer else None
|
||||
save_checkpoint(self.model, filepath, optimizer=optimizer, meta=meta)
|
||||
# in some environments, `os.symlink` is not supported, you may need to
|
||||
# set `create_symlink` to False
|
||||
# if create_symlink:
|
||||
# dst_file = osp.join(out_dir, 'latest.pth')
|
||||
# if platform.system() != 'Windows':
|
||||
# mmcv.symlink(filename, dst_file)
|
||||
# else:
|
||||
# shutil.copy(filepath, dst_file)
|
||||
|
||||
def resume(self,
|
||||
checkpoint,
|
||||
resume_optimizer=True,
|
||||
map_location='default'):
|
||||
if map_location == 'default':
|
||||
if torch.cuda.is_available():
|
||||
device_id = torch.cuda.current_device()
|
||||
checkpoint = self.load_checkpoint(
|
||||
checkpoint,
|
||||
map_location=lambda storage, loc: storage.cuda(device_id))
|
||||
else:
|
||||
checkpoint = self.load_checkpoint(checkpoint)
|
||||
else:
|
||||
checkpoint = self.load_checkpoint(
|
||||
checkpoint, map_location=map_location)
|
||||
|
||||
self._epoch = checkpoint['meta']['epoch']
|
||||
self._iter = checkpoint['meta']['iter']
|
||||
self._inner_iter = checkpoint['meta']['iter']
|
||||
if 'optimizer' in checkpoint and resume_optimizer:
|
||||
if isinstance(self.optimizer, Optimizer):
|
||||
self.optimizer.load_state_dict(checkpoint['optimizer'])
|
||||
elif isinstance(self.optimizer, dict):
|
||||
for k in self.optimizer.keys():
|
||||
self.optimizer[k].load_state_dict(
|
||||
checkpoint['optimizer'][k])
|
||||
else:
|
||||
raise TypeError(
|
||||
'Optimizer should be dict or torch.optim.Optimizer '
|
||||
f'but got {type(self.optimizer)}')
|
||||
|
||||
if 'amp' in checkpoint:
|
||||
apex.amp.load_state_dict(checkpoint['amp'])
|
||||
self.logger.info('load amp state dict')
|
||||
|
||||
self.logger.info(f'resumed from epoch: {self.epoch}, iter {self.iter}')
|
||||
@@ -0,0 +1,85 @@
|
||||
# Copyright (c) Open-MMLab. All rights reserved.
|
||||
import os.path as osp
|
||||
import time
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
import torch
|
||||
from torch.optim import Optimizer
|
||||
|
||||
import mmcv
|
||||
from mmcv.parallel import is_module_wrapper
|
||||
from mmcv.runner.checkpoint import weights_to_cpu, get_state_dict
|
||||
|
||||
try:
|
||||
import apex
|
||||
except:
|
||||
print('apex is not installed')
|
||||
|
||||
|
||||
def save_checkpoint(model, filename, optimizer=None, meta=None):
|
||||
"""Save checkpoint to file.
|
||||
|
||||
The checkpoint will have 4 fields: ``meta``, ``state_dict`` and
|
||||
``optimizer``, ``amp``. By default ``meta`` will contain version
|
||||
and time info.
|
||||
|
||||
Args:
|
||||
model (Module): Module whose params are to be saved.
|
||||
filename (str): Checkpoint filename.
|
||||
optimizer (:obj:`Optimizer`, optional): Optimizer to be saved.
|
||||
meta (dict, optional): Metadata to be saved in checkpoint.
|
||||
"""
|
||||
if meta is None:
|
||||
meta = {}
|
||||
elif not isinstance(meta, dict):
|
||||
raise TypeError(f'meta must be a dict or None, but got {type(meta)}')
|
||||
meta.update(mmcv_version=mmcv.__version__, time=time.asctime())
|
||||
|
||||
if is_module_wrapper(model):
|
||||
model = model.module
|
||||
|
||||
if hasattr(model, 'CLASSES') and model.CLASSES is not None:
|
||||
# save class name to the meta
|
||||
meta.update(CLASSES=model.CLASSES)
|
||||
|
||||
checkpoint = {
|
||||
'meta': meta,
|
||||
'state_dict': weights_to_cpu(get_state_dict(model))
|
||||
}
|
||||
# save optimizer state dict in the checkpoint
|
||||
if isinstance(optimizer, Optimizer):
|
||||
checkpoint['optimizer'] = optimizer.state_dict()
|
||||
elif isinstance(optimizer, dict):
|
||||
checkpoint['optimizer'] = {}
|
||||
for name, optim in optimizer.items():
|
||||
checkpoint['optimizer'][name] = optim.state_dict()
|
||||
|
||||
# save amp state dict in the checkpoint
|
||||
checkpoint['amp'] = apex.amp.state_dict()
|
||||
|
||||
if filename.startswith('pavi://'):
|
||||
try:
|
||||
from pavi import modelcloud
|
||||
from pavi.exception import NodeNotFoundError
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
'Please install pavi to load checkpoint from modelcloud.')
|
||||
model_path = filename[7:]
|
||||
root = modelcloud.Folder()
|
||||
model_dir, model_name = osp.split(model_path)
|
||||
try:
|
||||
model = modelcloud.get(model_dir)
|
||||
except NodeNotFoundError:
|
||||
model = root.create_training_model(model_dir)
|
||||
with TemporaryDirectory() as tmp_dir:
|
||||
checkpoint_file = osp.join(tmp_dir, model_name)
|
||||
with open(checkpoint_file, 'wb') as f:
|
||||
torch.save(checkpoint, f)
|
||||
f.flush()
|
||||
model.create_file(checkpoint_file, name=model_name)
|
||||
else:
|
||||
mmcv.mkdir_or_exist(osp.dirname(filename))
|
||||
# immediately flush buffer
|
||||
with open(filename, 'wb') as f:
|
||||
torch.save(checkpoint, f)
|
||||
f.flush()
|
||||
@@ -0,0 +1,33 @@
|
||||
from mmcv.runner import OptimizerHook, HOOKS
|
||||
try:
|
||||
import apex
|
||||
except:
|
||||
print('apex is not installed')
|
||||
|
||||
|
||||
@HOOKS.register_module()
|
||||
class DistOptimizerHook_custom(OptimizerHook):
|
||||
"""Optimizer hook for distributed training."""
|
||||
|
||||
def __init__(self, update_interval=1, grad_clip=None, coalesce=True, bucket_size_mb=-1, use_fp16=False):
|
||||
self.grad_clip = grad_clip
|
||||
self.coalesce = coalesce
|
||||
self.bucket_size_mb = bucket_size_mb
|
||||
self.update_interval = update_interval
|
||||
self.use_fp16 = use_fp16
|
||||
|
||||
def before_run(self, runner):
|
||||
runner.optimizer.zero_grad()
|
||||
|
||||
def after_train_iter(self, runner):
|
||||
runner.outputs['loss'] /= self.update_interval
|
||||
if self.use_fp16:
|
||||
with apex.amp.scale_loss(runner.outputs['loss'], runner.optimizer) as scaled_loss:
|
||||
scaled_loss.backward()
|
||||
else:
|
||||
runner.outputs['loss'].backward()
|
||||
if self.every_n_iters(runner, self.update_interval):
|
||||
if self.grad_clip is not None:
|
||||
self.clip_grads(runner.model.parameters())
|
||||
runner.optimizer.step()
|
||||
runner.optimizer.zero_grad()
|
||||
Reference in New Issue
Block a user