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,24 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
GRAD_CLIP = .01
|
||||
|
||||
class GradClip(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, x):
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_x):
|
||||
o = torch.zeros_like(grad_x)
|
||||
grad_x = torch.where(grad_x.abs()>GRAD_CLIP, o, grad_x)
|
||||
grad_x = torch.where(torch.isnan(grad_x), o, grad_x)
|
||||
return grad_x
|
||||
|
||||
class GradientClip(nn.Module):
|
||||
def __init__(self):
|
||||
super(GradientClip, self).__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return GradClip.apply(x)
|
||||
@@ -0,0 +1,140 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import droid_backends
|
||||
|
||||
class CorrSampler(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, volume, coords, radius):
|
||||
ctx.save_for_backward(volume,coords)
|
||||
ctx.radius = radius
|
||||
corr, = droid_backends.corr_index_forward(volume, coords, radius)
|
||||
return corr
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
volume, coords = ctx.saved_tensors
|
||||
grad_output = grad_output.contiguous()
|
||||
grad_volume, = droid_backends.corr_index_backward(volume, coords, grad_output, ctx.radius)
|
||||
return grad_volume, None, None
|
||||
|
||||
|
||||
class CorrBlock:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=3):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
self.corr_pyramid = []
|
||||
|
||||
# all pairs correlation
|
||||
corr = CorrBlock.corr(fmap1, fmap2)
|
||||
|
||||
batch, num, h1, w1, h2, w2 = corr.shape
|
||||
corr = corr.reshape(batch*num*h1*w1, 1, h2, w2)
|
||||
|
||||
for i in range(self.num_levels):
|
||||
self.corr_pyramid.append(
|
||||
corr.view(batch*num, h1, w1, h2//2**i, w2//2**i))
|
||||
corr = F.avg_pool2d(corr, 2, stride=2)
|
||||
|
||||
def __call__(self, coords):
|
||||
out_pyramid = []
|
||||
batch, num, ht, wd, _ = coords.shape
|
||||
coords = coords.permute(0,1,4,2,3)
|
||||
coords = coords.contiguous().view(batch*num, 2, ht, wd)
|
||||
|
||||
for i in range(self.num_levels):
|
||||
corr = CorrSampler.apply(self.corr_pyramid[i], coords/2**i, self.radius)
|
||||
out_pyramid.append(corr.view(batch, num, -1, ht, wd))
|
||||
|
||||
return torch.cat(out_pyramid, dim=2)
|
||||
|
||||
def cat(self, other):
|
||||
for i in range(self.num_levels):
|
||||
self.corr_pyramid[i] = torch.cat([self.corr_pyramid[i], other.corr_pyramid[i]], 0)
|
||||
return self
|
||||
|
||||
def __getitem__(self, index):
|
||||
for i in range(self.num_levels):
|
||||
self.corr_pyramid[i] = self.corr_pyramid[i][index]
|
||||
return self
|
||||
|
||||
|
||||
@staticmethod
|
||||
def corr(fmap1, fmap2):
|
||||
""" all-pairs correlation """
|
||||
batch, num, dim, ht, wd = fmap1.shape
|
||||
fmap1 = fmap1.reshape(batch*num, dim, ht*wd) / 4.0
|
||||
fmap2 = fmap2.reshape(batch*num, dim, ht*wd) / 4.0
|
||||
|
||||
corr = torch.matmul(fmap1.transpose(1,2), fmap2)
|
||||
return corr.view(batch, num, ht, wd, ht, wd)
|
||||
|
||||
|
||||
class CorrLayer(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, fmap1, fmap2, coords, r):
|
||||
ctx.r = r
|
||||
ctx.save_for_backward(fmap1, fmap2, coords)
|
||||
corr, = droid_backends.altcorr_forward(fmap1, fmap2, coords, ctx.r)
|
||||
return corr
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_corr):
|
||||
fmap1, fmap2, coords = ctx.saved_tensors
|
||||
grad_corr = grad_corr.contiguous()
|
||||
fmap1_grad, fmap2_grad, coords_grad = \
|
||||
droid_backends.altcorr_backward(fmap1, fmap2, coords, grad_corr, ctx.r)
|
||||
return fmap1_grad, fmap2_grad, coords_grad, None
|
||||
|
||||
|
||||
class AltCorrBlock:
|
||||
def __init__(self, fmaps, num_levels=4, radius=3):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
|
||||
B, N, C, H, W = fmaps.shape
|
||||
fmaps = fmaps.view(B*N, C, H, W) / 4.0
|
||||
|
||||
self.pyramid = []
|
||||
for i in range(self.num_levels):
|
||||
sz = (B, N, H//2**i, W//2**i, C)
|
||||
fmap_lvl = fmaps.permute(0, 2, 3, 1).contiguous()
|
||||
self.pyramid.append(fmap_lvl.view(*sz))
|
||||
fmaps = F.avg_pool2d(fmaps, 2, stride=2)
|
||||
|
||||
def corr_fn(self, coords, ii, jj):
|
||||
B, N, H, W, S, _ = coords.shape
|
||||
coords = coords.permute(0, 1, 4, 2, 3, 5)
|
||||
|
||||
corr_list = []
|
||||
for i in range(self.num_levels):
|
||||
r = self.radius
|
||||
fmap1_i = self.pyramid[0][:, ii]
|
||||
fmap2_i = self.pyramid[i][:, jj]
|
||||
|
||||
coords_i = (coords / 2**i).reshape(B*N, S, H, W, 2).contiguous()
|
||||
fmap1_i = fmap1_i.reshape((B*N,) + fmap1_i.shape[2:])
|
||||
fmap2_i = fmap2_i.reshape((B*N,) + fmap2_i.shape[2:])
|
||||
|
||||
corr = CorrLayer.apply(fmap1_i.float(), fmap2_i.float(), coords_i, self.radius)
|
||||
corr = corr.view(B, N, S, -1, H, W).permute(0, 1, 3, 4, 5, 2)
|
||||
corr_list.append(corr)
|
||||
|
||||
corr = torch.cat(corr_list, dim=2)
|
||||
return corr
|
||||
|
||||
|
||||
def __call__(self, coords, ii, jj):
|
||||
squeeze_output = False
|
||||
if len(coords.shape) == 5:
|
||||
coords = coords.unsqueeze(dim=-2)
|
||||
squeeze_output = True
|
||||
|
||||
corr = self.corr_fn(coords, ii, jj)
|
||||
|
||||
if squeeze_output:
|
||||
corr = corr.squeeze(dim=-1)
|
||||
|
||||
return corr.contiguous()
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn='group', stride=1):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(planes)
|
||||
self.norm2 = nn.BatchNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(planes)
|
||||
self.norm2 = nn.InstanceNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.Sequential()
|
||||
|
||||
if stride == 1:
|
||||
self.downsample = None
|
||||
|
||||
else:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.relu(self.norm1(self.conv1(y)))
|
||||
y = self.relu(self.norm2(self.conv2(y)))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x+y)
|
||||
|
||||
|
||||
class BottleneckBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn='group', stride=1):
|
||||
super(BottleneckBlock, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(in_planes, planes//4, kernel_size=1, padding=0)
|
||||
self.conv2 = nn.Conv2d(planes//4, planes//4, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv3 = nn.Conv2d(planes//4, planes, kernel_size=1, padding=0)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(planes//4)
|
||||
self.norm2 = nn.BatchNorm2d(planes//4)
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(planes//4)
|
||||
self.norm2 = nn.InstanceNorm2d(planes//4)
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
self.norm3 = nn.Sequential()
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.Sequential()
|
||||
|
||||
if stride == 1:
|
||||
self.downsample = None
|
||||
|
||||
else:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4)
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.relu(self.norm1(self.conv1(y)))
|
||||
y = self.relu(self.norm2(self.conv2(y)))
|
||||
y = self.relu(self.norm3(self.conv3(y)))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x+y)
|
||||
|
||||
|
||||
DIM=32
|
||||
|
||||
class BasicEncoder(nn.Module):
|
||||
def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0, multidim=False):
|
||||
super(BasicEncoder, self).__init__()
|
||||
self.norm_fn = norm_fn
|
||||
self.multidim = multidim
|
||||
|
||||
if self.norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=8, num_channels=DIM)
|
||||
|
||||
elif self.norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(DIM)
|
||||
|
||||
elif self.norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(DIM)
|
||||
|
||||
elif self.norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
|
||||
self.conv1 = nn.Conv2d(3, DIM, kernel_size=7, stride=2, padding=3)
|
||||
self.relu1 = nn.ReLU(inplace=True)
|
||||
|
||||
self.in_planes = DIM
|
||||
self.layer1 = self._make_layer(DIM, stride=1)
|
||||
self.layer2 = self._make_layer(2*DIM, stride=2)
|
||||
self.layer3 = self._make_layer(4*DIM, stride=2)
|
||||
|
||||
# output convolution
|
||||
self.conv2 = nn.Conv2d(4*DIM, output_dim, kernel_size=1)
|
||||
|
||||
if self.multidim:
|
||||
self.layer4 = self._make_layer(256, stride=2)
|
||||
self.layer5 = self._make_layer(512, stride=2)
|
||||
|
||||
self.in_planes = 256
|
||||
self.layer6 = self._make_layer(256, stride=1)
|
||||
|
||||
self.in_planes = 128
|
||||
self.layer7 = self._make_layer(128, stride=1)
|
||||
|
||||
self.up1 = nn.Conv2d(512, 256, 1)
|
||||
self.up2 = nn.Conv2d(256, 128, 1)
|
||||
self.conv3 = nn.Conv2d(128, output_dim, kernel_size=1)
|
||||
|
||||
if dropout > 0:
|
||||
self.dropout = nn.Dropout2d(p=dropout)
|
||||
else:
|
||||
self.dropout = None
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
|
||||
if m.weight is not None:
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_layer(self, dim, stride=1):
|
||||
layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
|
||||
layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
|
||||
layers = (layer1, layer2)
|
||||
|
||||
self.in_planes = dim
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
b, n, c1, h1, w1 = x.shape
|
||||
x = x.view(b*n, c1, h1, w1)
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.relu1(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
|
||||
x = self.conv2(x)
|
||||
|
||||
_, c2, h2, w2 = x.shape
|
||||
return x.view(b, n, c2, h2, w2)
|
||||
@@ -0,0 +1,34 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class ConvGRU(nn.Module):
|
||||
def __init__(self, h_planes=128, i_planes=128):
|
||||
super(ConvGRU, self).__init__()
|
||||
self.do_checkpoint = False
|
||||
self.convz = nn.Conv2d(h_planes+i_planes, h_planes, 3, padding=1)
|
||||
self.convr = nn.Conv2d(h_planes+i_planes, h_planes, 3, padding=1)
|
||||
self.convq = nn.Conv2d(h_planes+i_planes, h_planes, 3, padding=1)
|
||||
|
||||
self.w = nn.Conv2d(h_planes, h_planes, 1, padding=0)
|
||||
|
||||
self.convz_glo = nn.Conv2d(h_planes, h_planes, 1, padding=0)
|
||||
self.convr_glo = nn.Conv2d(h_planes, h_planes, 1, padding=0)
|
||||
self.convq_glo = nn.Conv2d(h_planes, h_planes, 1, padding=0)
|
||||
|
||||
def forward(self, net, *inputs):
|
||||
inp = torch.cat(inputs, dim=1)
|
||||
net_inp = torch.cat([net, inp], dim=1)
|
||||
|
||||
b, c, h, w = net.shape
|
||||
glo = torch.sigmoid(self.w(net)) * net
|
||||
glo = glo.view(b, c, h*w).mean(-1).view(b, c, 1, 1)
|
||||
|
||||
z = torch.sigmoid(self.convz(net_inp) + self.convz_glo(glo))
|
||||
r = torch.sigmoid(self.convr(net_inp) + self.convr_glo(glo))
|
||||
q = torch.tanh(self.convq(torch.cat([r*net, inp], dim=1)) + self.convq_glo(glo))
|
||||
|
||||
net = (1-z) * net + z * q
|
||||
return net
|
||||
|
||||
|
||||
Reference in New Issue
Block a user