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,616 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import copy as cp
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from mmcv.cnn import (ConvModule, MaxPool2d, constant_init, kaiming_init,
|
||||
normal_init)
|
||||
|
||||
from ..builder import BACKBONES
|
||||
from .base_backbone import BaseBackbone
|
||||
|
||||
|
||||
class RSB(nn.Module):
|
||||
"""Residual Steps block for RSN. Paper ref: Cai et al. "Learning Delicate
|
||||
Local Representations for Multi-Person Pose Estimation" (ECCV 2020).
|
||||
|
||||
Args:
|
||||
in_channels (int): Input channels of this block.
|
||||
out_channels (int): Output channels of this block.
|
||||
num_steps (int): Numbers of steps in RSB
|
||||
stride (int): stride of the block. Default: 1
|
||||
downsample (nn.Module): downsample operation on identity branch.
|
||||
Default: None.
|
||||
norm_cfg (dict): dictionary to construct and config norm layer.
|
||||
Default: dict(type='BN')
|
||||
expand_times (int): Times by which the in_channels are expanded.
|
||||
Default:26.
|
||||
res_top_channels (int): Number of channels of feature output by
|
||||
ResNet_top. Default:64.
|
||||
"""
|
||||
|
||||
expansion = 1
|
||||
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
num_steps=4,
|
||||
stride=1,
|
||||
downsample=None,
|
||||
with_cp=False,
|
||||
norm_cfg=dict(type='BN'),
|
||||
expand_times=26,
|
||||
res_top_channels=64):
|
||||
# Protect mutable default arguments
|
||||
norm_cfg = cp.deepcopy(norm_cfg)
|
||||
super().__init__()
|
||||
assert num_steps > 1
|
||||
self.in_channels = in_channels
|
||||
self.branch_channels = self.in_channels * expand_times
|
||||
self.branch_channels //= res_top_channels
|
||||
self.out_channels = out_channels
|
||||
self.stride = stride
|
||||
self.downsample = downsample
|
||||
self.with_cp = with_cp
|
||||
self.norm_cfg = norm_cfg
|
||||
self.num_steps = num_steps
|
||||
self.conv_bn_relu1 = ConvModule(
|
||||
self.in_channels,
|
||||
self.num_steps * self.branch_channels,
|
||||
kernel_size=1,
|
||||
stride=self.stride,
|
||||
padding=0,
|
||||
norm_cfg=self.norm_cfg,
|
||||
inplace=False)
|
||||
for i in range(self.num_steps):
|
||||
for j in range(i + 1):
|
||||
module_name = f'conv_bn_relu2_{i + 1}_{j + 1}'
|
||||
self.add_module(
|
||||
module_name,
|
||||
ConvModule(
|
||||
self.branch_channels,
|
||||
self.branch_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
norm_cfg=self.norm_cfg,
|
||||
inplace=False))
|
||||
self.conv_bn3 = ConvModule(
|
||||
self.num_steps * self.branch_channels,
|
||||
self.out_channels * self.expansion,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
act_cfg=None,
|
||||
norm_cfg=self.norm_cfg,
|
||||
inplace=False)
|
||||
self.relu = nn.ReLU(inplace=False)
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward function."""
|
||||
|
||||
identity = x
|
||||
x = self.conv_bn_relu1(x)
|
||||
spx = torch.split(x, self.branch_channels, 1)
|
||||
outputs = list()
|
||||
outs = list()
|
||||
for i in range(self.num_steps):
|
||||
outputs_i = list()
|
||||
outputs.append(outputs_i)
|
||||
for j in range(i + 1):
|
||||
if j == 0:
|
||||
inputs = spx[i]
|
||||
else:
|
||||
inputs = outputs[i][j - 1]
|
||||
if i > j:
|
||||
inputs = inputs + outputs[i - 1][j]
|
||||
module_name = f'conv_bn_relu2_{i + 1}_{j + 1}'
|
||||
module_i_j = getattr(self, module_name)
|
||||
outputs[i].append(module_i_j(inputs))
|
||||
|
||||
outs.append(outputs[i][i])
|
||||
out = torch.cat(tuple(outs), 1)
|
||||
out = self.conv_bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(identity)
|
||||
out = out + identity
|
||||
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Downsample_module(nn.Module):
|
||||
"""Downsample module for RSN.
|
||||
|
||||
Args:
|
||||
block (nn.Module): Downsample block.
|
||||
num_blocks (list): Number of blocks in each downsample unit.
|
||||
num_units (int): Numbers of downsample units. Default: 4
|
||||
has_skip (bool): Have skip connections from prior upsample
|
||||
module or not. Default:False
|
||||
num_steps (int): Number of steps in a block. Default:4
|
||||
norm_cfg (dict): dictionary to construct and config norm layer.
|
||||
Default: dict(type='BN')
|
||||
in_channels (int): Number of channels of the input feature to
|
||||
downsample module. Default: 64
|
||||
expand_times (int): Times by which the in_channels are expanded.
|
||||
Default:26.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
block,
|
||||
num_blocks,
|
||||
num_steps=4,
|
||||
num_units=4,
|
||||
has_skip=False,
|
||||
norm_cfg=dict(type='BN'),
|
||||
in_channels=64,
|
||||
expand_times=26):
|
||||
# Protect mutable default arguments
|
||||
norm_cfg = cp.deepcopy(norm_cfg)
|
||||
super().__init__()
|
||||
self.has_skip = has_skip
|
||||
self.in_channels = in_channels
|
||||
assert len(num_blocks) == num_units
|
||||
self.num_blocks = num_blocks
|
||||
self.num_units = num_units
|
||||
self.num_steps = num_steps
|
||||
self.norm_cfg = norm_cfg
|
||||
self.layer1 = self._make_layer(
|
||||
block,
|
||||
in_channels,
|
||||
num_blocks[0],
|
||||
expand_times=expand_times,
|
||||
res_top_channels=in_channels)
|
||||
for i in range(1, num_units):
|
||||
module_name = f'layer{i + 1}'
|
||||
self.add_module(
|
||||
module_name,
|
||||
self._make_layer(
|
||||
block,
|
||||
in_channels * pow(2, i),
|
||||
num_blocks[i],
|
||||
stride=2,
|
||||
expand_times=expand_times,
|
||||
res_top_channels=in_channels))
|
||||
|
||||
def _make_layer(self,
|
||||
block,
|
||||
out_channels,
|
||||
blocks,
|
||||
stride=1,
|
||||
expand_times=26,
|
||||
res_top_channels=64):
|
||||
downsample = None
|
||||
if stride != 1 or self.in_channels != out_channels * block.expansion:
|
||||
downsample = ConvModule(
|
||||
self.in_channels,
|
||||
out_channels * block.expansion,
|
||||
kernel_size=1,
|
||||
stride=stride,
|
||||
padding=0,
|
||||
norm_cfg=self.norm_cfg,
|
||||
act_cfg=None,
|
||||
inplace=True)
|
||||
|
||||
units = list()
|
||||
units.append(
|
||||
block(
|
||||
self.in_channels,
|
||||
out_channels,
|
||||
num_steps=self.num_steps,
|
||||
stride=stride,
|
||||
downsample=downsample,
|
||||
norm_cfg=self.norm_cfg,
|
||||
expand_times=expand_times,
|
||||
res_top_channels=res_top_channels))
|
||||
self.in_channels = out_channels * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
units.append(
|
||||
block(
|
||||
self.in_channels,
|
||||
out_channels,
|
||||
num_steps=self.num_steps,
|
||||
expand_times=expand_times,
|
||||
res_top_channels=res_top_channels))
|
||||
|
||||
return nn.Sequential(*units)
|
||||
|
||||
def forward(self, x, skip1, skip2):
|
||||
out = list()
|
||||
for i in range(self.num_units):
|
||||
module_name = f'layer{i + 1}'
|
||||
module_i = getattr(self, module_name)
|
||||
x = module_i(x)
|
||||
if self.has_skip:
|
||||
x = x + skip1[i] + skip2[i]
|
||||
out.append(x)
|
||||
out.reverse()
|
||||
|
||||
return tuple(out)
|
||||
|
||||
|
||||
class Upsample_unit(nn.Module):
|
||||
"""Upsample unit for upsample module.
|
||||
|
||||
Args:
|
||||
ind (int): Indicates whether to interpolate (>0) and whether to
|
||||
generate feature map for the next hourglass-like module.
|
||||
num_units (int): Number of units that form a upsample module. Along
|
||||
with ind and gen_cross_conv, nm_units is used to decide whether
|
||||
to generate feature map for the next hourglass-like module.
|
||||
in_channels (int): Channel number of the skip-in feature maps from
|
||||
the corresponding downsample unit.
|
||||
unit_channels (int): Channel number in this unit. Default:256.
|
||||
gen_skip: (bool): Whether or not to generate skips for the posterior
|
||||
downsample module. Default:False
|
||||
gen_cross_conv (bool): Whether to generate feature map for the next
|
||||
hourglass-like module. Default:False
|
||||
norm_cfg (dict): dictionary to construct and config norm layer.
|
||||
Default: dict(type='BN')
|
||||
out_channels (in): Number of channels of feature output by upsample
|
||||
module. Must equal to in_channels of downsample module. Default:64
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
ind,
|
||||
num_units,
|
||||
in_channels,
|
||||
unit_channels=256,
|
||||
gen_skip=False,
|
||||
gen_cross_conv=False,
|
||||
norm_cfg=dict(type='BN'),
|
||||
out_channels=64):
|
||||
# Protect mutable default arguments
|
||||
norm_cfg = cp.deepcopy(norm_cfg)
|
||||
super().__init__()
|
||||
self.num_units = num_units
|
||||
self.norm_cfg = norm_cfg
|
||||
self.in_skip = ConvModule(
|
||||
in_channels,
|
||||
unit_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
norm_cfg=self.norm_cfg,
|
||||
act_cfg=None,
|
||||
inplace=True)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
self.ind = ind
|
||||
if self.ind > 0:
|
||||
self.up_conv = ConvModule(
|
||||
unit_channels,
|
||||
unit_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
norm_cfg=self.norm_cfg,
|
||||
act_cfg=None,
|
||||
inplace=True)
|
||||
|
||||
self.gen_skip = gen_skip
|
||||
if self.gen_skip:
|
||||
self.out_skip1 = ConvModule(
|
||||
in_channels,
|
||||
in_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
norm_cfg=self.norm_cfg,
|
||||
inplace=True)
|
||||
|
||||
self.out_skip2 = ConvModule(
|
||||
unit_channels,
|
||||
in_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
norm_cfg=self.norm_cfg,
|
||||
inplace=True)
|
||||
|
||||
self.gen_cross_conv = gen_cross_conv
|
||||
if self.ind == num_units - 1 and self.gen_cross_conv:
|
||||
self.cross_conv = ConvModule(
|
||||
unit_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
norm_cfg=self.norm_cfg,
|
||||
inplace=True)
|
||||
|
||||
def forward(self, x, up_x):
|
||||
out = self.in_skip(x)
|
||||
|
||||
if self.ind > 0:
|
||||
up_x = F.interpolate(
|
||||
up_x,
|
||||
size=(x.size(2), x.size(3)),
|
||||
mode='bilinear',
|
||||
align_corners=True)
|
||||
up_x = self.up_conv(up_x)
|
||||
out = out + up_x
|
||||
out = self.relu(out)
|
||||
|
||||
skip1 = None
|
||||
skip2 = None
|
||||
if self.gen_skip:
|
||||
skip1 = self.out_skip1(x)
|
||||
skip2 = self.out_skip2(out)
|
||||
|
||||
cross_conv = None
|
||||
if self.ind == self.num_units - 1 and self.gen_cross_conv:
|
||||
cross_conv = self.cross_conv(out)
|
||||
|
||||
return out, skip1, skip2, cross_conv
|
||||
|
||||
|
||||
class Upsample_module(nn.Module):
|
||||
"""Upsample module for RSN.
|
||||
|
||||
Args:
|
||||
unit_channels (int): Channel number in the upsample units.
|
||||
Default:256.
|
||||
num_units (int): Numbers of upsample units. Default: 4
|
||||
gen_skip (bool): Whether to generate skip for posterior downsample
|
||||
module or not. Default:False
|
||||
gen_cross_conv (bool): Whether to generate feature map for the next
|
||||
hourglass-like module. Default:False
|
||||
norm_cfg (dict): dictionary to construct and config norm layer.
|
||||
Default: dict(type='BN')
|
||||
out_channels (int): Number of channels of feature output by upsample
|
||||
module. Must equal to in_channels of downsample module. Default:64
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
unit_channels=256,
|
||||
num_units=4,
|
||||
gen_skip=False,
|
||||
gen_cross_conv=False,
|
||||
norm_cfg=dict(type='BN'),
|
||||
out_channels=64):
|
||||
# Protect mutable default arguments
|
||||
norm_cfg = cp.deepcopy(norm_cfg)
|
||||
super().__init__()
|
||||
self.in_channels = list()
|
||||
for i in range(num_units):
|
||||
self.in_channels.append(RSB.expansion * out_channels * pow(2, i))
|
||||
self.in_channels.reverse()
|
||||
self.num_units = num_units
|
||||
self.gen_skip = gen_skip
|
||||
self.gen_cross_conv = gen_cross_conv
|
||||
self.norm_cfg = norm_cfg
|
||||
for i in range(num_units):
|
||||
module_name = f'up{i + 1}'
|
||||
self.add_module(
|
||||
module_name,
|
||||
Upsample_unit(
|
||||
i,
|
||||
self.num_units,
|
||||
self.in_channels[i],
|
||||
unit_channels,
|
||||
self.gen_skip,
|
||||
self.gen_cross_conv,
|
||||
norm_cfg=self.norm_cfg,
|
||||
out_channels=64))
|
||||
|
||||
def forward(self, x):
|
||||
out = list()
|
||||
skip1 = list()
|
||||
skip2 = list()
|
||||
cross_conv = None
|
||||
for i in range(self.num_units):
|
||||
module_i = getattr(self, f'up{i + 1}')
|
||||
if i == 0:
|
||||
outi, skip1_i, skip2_i, _ = module_i(x[i], None)
|
||||
elif i == self.num_units - 1:
|
||||
outi, skip1_i, skip2_i, cross_conv = module_i(x[i], out[i - 1])
|
||||
else:
|
||||
outi, skip1_i, skip2_i, _ = module_i(x[i], out[i - 1])
|
||||
out.append(outi)
|
||||
skip1.append(skip1_i)
|
||||
skip2.append(skip2_i)
|
||||
skip1.reverse()
|
||||
skip2.reverse()
|
||||
|
||||
return out, skip1, skip2, cross_conv
|
||||
|
||||
|
||||
class Single_stage_RSN(nn.Module):
|
||||
"""Single_stage Residual Steps Network.
|
||||
|
||||
Args:
|
||||
unit_channels (int): Channel number in the upsample units. Default:256.
|
||||
num_units (int): Numbers of downsample/upsample units. Default: 4
|
||||
gen_skip (bool): Whether to generate skip for posterior downsample
|
||||
module or not. Default:False
|
||||
gen_cross_conv (bool): Whether to generate feature map for the next
|
||||
hourglass-like module. Default:False
|
||||
has_skip (bool): Have skip connections from prior upsample
|
||||
module or not. Default:False
|
||||
num_steps (int): Number of steps in RSB. Default: 4
|
||||
num_blocks (list): Number of blocks in each downsample unit.
|
||||
Default: [2, 2, 2, 2] Note: Make sure num_units==len(num_blocks)
|
||||
norm_cfg (dict): dictionary to construct and config norm layer.
|
||||
Default: dict(type='BN')
|
||||
in_channels (int): Number of channels of the feature from ResNet_Top.
|
||||
Default: 64.
|
||||
expand_times (int): Times by which the in_channels are expanded in RSB.
|
||||
Default:26.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
has_skip=False,
|
||||
gen_skip=False,
|
||||
gen_cross_conv=False,
|
||||
unit_channels=256,
|
||||
num_units=4,
|
||||
num_steps=4,
|
||||
num_blocks=[2, 2, 2, 2],
|
||||
norm_cfg=dict(type='BN'),
|
||||
in_channels=64,
|
||||
expand_times=26):
|
||||
# Protect mutable default arguments
|
||||
norm_cfg = cp.deepcopy(norm_cfg)
|
||||
num_blocks = cp.deepcopy(num_blocks)
|
||||
super().__init__()
|
||||
assert len(num_blocks) == num_units
|
||||
self.has_skip = has_skip
|
||||
self.gen_skip = gen_skip
|
||||
self.gen_cross_conv = gen_cross_conv
|
||||
self.num_units = num_units
|
||||
self.num_steps = num_steps
|
||||
self.unit_channels = unit_channels
|
||||
self.num_blocks = num_blocks
|
||||
self.norm_cfg = norm_cfg
|
||||
|
||||
self.downsample = Downsample_module(RSB, num_blocks, num_steps,
|
||||
num_units, has_skip, norm_cfg,
|
||||
in_channels, expand_times)
|
||||
self.upsample = Upsample_module(unit_channels, num_units, gen_skip,
|
||||
gen_cross_conv, norm_cfg, in_channels)
|
||||
|
||||
def forward(self, x, skip1, skip2):
|
||||
mid = self.downsample(x, skip1, skip2)
|
||||
out, skip1, skip2, cross_conv = self.upsample(mid)
|
||||
|
||||
return out, skip1, skip2, cross_conv
|
||||
|
||||
|
||||
class ResNet_top(nn.Module):
|
||||
"""ResNet top for RSN.
|
||||
|
||||
Args:
|
||||
norm_cfg (dict): dictionary to construct and config norm layer.
|
||||
Default: dict(type='BN')
|
||||
channels (int): Number of channels of the feature output by ResNet_top.
|
||||
"""
|
||||
|
||||
def __init__(self, norm_cfg=dict(type='BN'), channels=64):
|
||||
# Protect mutable default arguments
|
||||
norm_cfg = cp.deepcopy(norm_cfg)
|
||||
super().__init__()
|
||||
self.top = nn.Sequential(
|
||||
ConvModule(
|
||||
3,
|
||||
channels,
|
||||
kernel_size=7,
|
||||
stride=2,
|
||||
padding=3,
|
||||
norm_cfg=norm_cfg,
|
||||
inplace=True), MaxPool2d(kernel_size=3, stride=2, padding=1))
|
||||
|
||||
def forward(self, img):
|
||||
return self.top(img)
|
||||
|
||||
|
||||
@BACKBONES.register_module()
|
||||
class RSN(BaseBackbone):
|
||||
"""Residual Steps Network backbone. Paper ref: Cai et al. "Learning
|
||||
Delicate Local Representations for Multi-Person Pose Estimation" (ECCV
|
||||
2020).
|
||||
|
||||
Args:
|
||||
unit_channels (int): Number of Channels in an upsample unit.
|
||||
Default: 256
|
||||
num_stages (int): Number of stages in a multi-stage RSN. Default: 4
|
||||
num_units (int): NUmber of downsample/upsample units in a single-stage
|
||||
RSN. Default: 4 Note: Make sure num_units == len(self.num_blocks)
|
||||
num_blocks (list): Number of RSBs (Residual Steps Block) in each
|
||||
downsample unit. Default: [2, 2, 2, 2]
|
||||
num_steps (int): Number of steps in a RSB. Default:4
|
||||
norm_cfg (dict): dictionary to construct and config norm layer.
|
||||
Default: dict(type='BN')
|
||||
res_top_channels (int): Number of channels of feature from ResNet_top.
|
||||
Default: 64.
|
||||
expand_times (int): Times by which the in_channels are expanded in RSB.
|
||||
Default:26.
|
||||
Example:
|
||||
>>> from mmpose.models import RSN
|
||||
>>> import torch
|
||||
>>> self = RSN(num_stages=2,num_units=2,num_blocks=[2,2])
|
||||
>>> self.eval()
|
||||
>>> inputs = torch.rand(1, 3, 511, 511)
|
||||
>>> level_outputs = self.forward(inputs)
|
||||
>>> for level_output in level_outputs:
|
||||
... for feature in level_output:
|
||||
... print(tuple(feature.shape))
|
||||
...
|
||||
(1, 256, 64, 64)
|
||||
(1, 256, 128, 128)
|
||||
(1, 256, 64, 64)
|
||||
(1, 256, 128, 128)
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
unit_channels=256,
|
||||
num_stages=4,
|
||||
num_units=4,
|
||||
num_blocks=[2, 2, 2, 2],
|
||||
num_steps=4,
|
||||
norm_cfg=dict(type='BN'),
|
||||
res_top_channels=64,
|
||||
expand_times=26):
|
||||
# Protect mutable default arguments
|
||||
norm_cfg = cp.deepcopy(norm_cfg)
|
||||
num_blocks = cp.deepcopy(num_blocks)
|
||||
super().__init__()
|
||||
self.unit_channels = unit_channels
|
||||
self.num_stages = num_stages
|
||||
self.num_units = num_units
|
||||
self.num_blocks = num_blocks
|
||||
self.num_steps = num_steps
|
||||
self.norm_cfg = norm_cfg
|
||||
|
||||
assert self.num_stages > 0
|
||||
assert self.num_steps > 1
|
||||
assert self.num_units > 1
|
||||
assert self.num_units == len(self.num_blocks)
|
||||
self.top = ResNet_top(norm_cfg=norm_cfg)
|
||||
self.multi_stage_rsn = nn.ModuleList([])
|
||||
for i in range(self.num_stages):
|
||||
if i == 0:
|
||||
has_skip = False
|
||||
else:
|
||||
has_skip = True
|
||||
if i != self.num_stages - 1:
|
||||
gen_skip = True
|
||||
gen_cross_conv = True
|
||||
else:
|
||||
gen_skip = False
|
||||
gen_cross_conv = False
|
||||
self.multi_stage_rsn.append(
|
||||
Single_stage_RSN(has_skip, gen_skip, gen_cross_conv,
|
||||
unit_channels, num_units, num_steps,
|
||||
num_blocks, norm_cfg, res_top_channels,
|
||||
expand_times))
|
||||
|
||||
def forward(self, x):
|
||||
"""Model forward function."""
|
||||
out_feats = []
|
||||
skip1 = None
|
||||
skip2 = None
|
||||
x = self.top(x)
|
||||
for i in range(self.num_stages):
|
||||
out, skip1, skip2, x = self.multi_stage_rsn[i](x, skip1, skip2)
|
||||
out_feats.append(out)
|
||||
|
||||
return out_feats
|
||||
|
||||
def init_weights(self, pretrained=None):
|
||||
"""Initialize model weights."""
|
||||
for m in self.multi_stage_rsn.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
kaiming_init(m)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
constant_init(m, 1)
|
||||
elif isinstance(m, nn.Linear):
|
||||
normal_init(m, std=0.01)
|
||||
|
||||
for m in self.top.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
kaiming_init(m)
|
||||
Reference in New Issue
Block a user