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,4 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .distributed_sampler import DistributedSampler
|
||||
|
||||
__all__ = ['DistributedSampler']
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
from torch.utils.data import DistributedSampler as _DistributedSampler
|
||||
|
||||
|
||||
class DistributedSampler(_DistributedSampler):
|
||||
"""DistributedSampler inheriting from
|
||||
`torch.utils.data.DistributedSampler`.
|
||||
|
||||
In pytorch of lower versions, there is no `shuffle` argument. This child
|
||||
class will port one to DistributedSampler.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
dataset,
|
||||
num_replicas=None,
|
||||
rank=None,
|
||||
shuffle=True,
|
||||
seed=0):
|
||||
super().__init__(
|
||||
dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)
|
||||
# for the compatibility from PyTorch 1.3+
|
||||
self.seed = seed if seed is not None else 0
|
||||
|
||||
def __iter__(self):
|
||||
"""Deterministically shuffle based on epoch."""
|
||||
if self.shuffle:
|
||||
g = torch.Generator()
|
||||
g.manual_seed(self.epoch + self.seed)
|
||||
indices = torch.randperm(len(self.dataset), generator=g).tolist()
|
||||
else:
|
||||
indices = torch.arange(len(self.dataset)).tolist()
|
||||
|
||||
# add extra samples to make it evenly divisible
|
||||
indices += indices[:(self.total_size - len(indices))]
|
||||
assert len(indices) == self.total_size
|
||||
|
||||
# subsample
|
||||
indices = indices[self.rank:self.total_size:self.num_replicas]
|
||||
assert len(indices) == self.num_samples
|
||||
return iter(indices)
|
||||
Reference in New Issue
Block a user