Skip to content

Commit

Permalink
Update the network definition to support models built with instance-n…
Browse files Browse the repository at this point in the history 
…orm instead of batch-norm
  • Loading branch information
@hhj1897
hhj1897 committed Dec 17, 2021
1 parent f5c6acc commit fa410a8
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 18 deletions.
44 changes: 29 additions & 15 deletions ibug/face_alignment/fan/fan.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,17 +9,21 @@ def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False):


class ConvBlock(nn.Module):
def __init__(self, in_planes, out_planes):
def __init__(self, in_planes, out_planes, use_instance_norm):
super(ConvBlock, self).__init__()
self.bn1 = nn.BatchNorm2d(in_planes)
self.bn1 = nn.InstanceNorm2d(in_planes) if use_instance_norm else nn.BatchNorm2d(in_planes)
self.conv1 = conv3x3(in_planes, int(out_planes / 2))
self.bn2 = nn.BatchNorm2d(int(out_planes / 2))
self.bn2 = (nn.InstanceNorm2d(int(out_planes / 2)) if use_instance_norm
else nn.BatchNorm2d(int(out_planes / 2)))
self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4))
self.bn3 = nn.BatchNorm2d(int(out_planes / 4))
self.bn3 = (nn.InstanceNorm2d(int(out_planes / 4)) if use_instance_norm
else nn.BatchNorm2d(int(out_planes / 4)))
self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4))

if in_planes != out_planes:
self.downsample = nn.Sequential(nn.BatchNorm2d(in_planes), nn.ReLU(True),
self.downsample = nn.Sequential(nn.InstanceNorm2d(in_planes) if use_instance_norm
else nn.BatchNorm2d(in_planes),
nn.ReLU(True),
nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, bias=False))
else:
self.downsample = None
Expand Down Expand Up @@ -57,17 +61,24 @@ def __init__(self, config):
self._generate_network(self.config.hg_depth)

def _generate_network(self, level):
self.add_module('b1_' + str(level), ConvBlock(self.config.hg_num_features, self.config.hg_num_features))
self.add_module('b1_' + str(level), ConvBlock(self.config.hg_num_features,
self.config.hg_num_features,
self.config.use_instance_norm))

self.add_module('b2_' + str(level), ConvBlock(self.config.hg_num_features, self.config.hg_num_features))
self.add_module('b2_' + str(level), ConvBlock(self.config.hg_num_features,
self.config.hg_num_features,
self.config.use_instance_norm))

if level > 1:
self._generate_network(level - 1)
else:
self.add_module('b2_plus_' + str(level),ConvBlock(self.config.hg_num_features,
self.config.hg_num_features))
self.config.hg_num_features,
self.config.use_instance_norm))

self.add_module('b3_' + str(level), ConvBlock(self.config.hg_num_features, self.config.hg_num_features))
self.add_module('b3_' + str(level), ConvBlock(self.config.hg_num_features,
self.config.hg_num_features,
self.config.use_instance_norm))

def _forward(self, level, inp):
up1 = inp
Expand Down Expand Up @@ -103,20 +114,23 @@ def __init__(self, config):

# Stem
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
self.bn1 = nn.BatchNorm2d(64)
self.conv2 = ConvBlock(64, 128)
self.conv3 = ConvBlock(128, 128)
self.conv4 = ConvBlock(128, self.config.hg_num_features)
self.bn1 = nn.InstanceNorm2d(64) if self.config.use_instance_norm else nn.BatchNorm2d(64)
self.conv2 = ConvBlock(64, 128, self.config.use_instance_norm)
self.conv3 = ConvBlock(128, 128, self.config.use_instance_norm)
self.conv4 = ConvBlock(128, self.config.hg_num_features, self.config.use_instance_norm)

# Hourglasses
for hg_module in range(self.config.num_modules):
self.add_module('m' + str(hg_module), HourGlass(self.config))
self.add_module('top_m_' + str(hg_module), ConvBlock(self.config.hg_num_features,
self.config.hg_num_features))
self.config.hg_num_features,
self.config.use_instance_norm))
self.add_module('conv_last' + str(hg_module), nn.Conv2d(self.config.hg_num_features,
self.config.hg_num_features,
kernel_size=1, stride=1, padding=0))
self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(self.config.hg_num_features))
self.add_module('bn_end' + str(hg_module),
nn.InstanceNorm2d(self.config.hg_num_features) if self.config.use_instance_norm
else nn.BatchNorm2d(self.config.hg_num_features))
self.add_module('l' + str(hg_module), nn.Conv2d(self.config.hg_num_features, 68,
kernel_size=1, stride=1, padding=0))

Expand Down
9 changes: 6 additions & 3 deletions ibug/face_alignment/fan/fan_predictor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,15 +32,18 @@ def get_model(name: str = '2dfan2') -> SimpleNamespace:
if name == '2dfan2':
return SimpleNamespace(weights=os.path.join(os.path.dirname(__file__), 'weights', '2dfan2.pth'),
config=SimpleNamespace(crop_ratio=0.55, input_size=256, num_modules=2,
hg_num_features=256, hg_depth=4, use_avg_pool=False))
hg_num_features=256, hg_depth=4, use_avg_pool=False,
use_instance_norm=False))
elif name == '2dfan4':
return SimpleNamespace(weights=os.path.join(os.path.dirname(__file__), 'weights', '2dfan4.pth'),
config=SimpleNamespace(crop_ratio=0.55, input_size=256, num_modules=4,
hg_num_features=256, hg_depth=4, use_avg_pool=True))
hg_num_features=256, hg_depth=4, use_avg_pool=True,
use_instance_norm=False))
elif name == '2dfan2_alt':
return SimpleNamespace(weights=os.path.join(os.path.dirname(__file__), 'weights', '2dfan2_alt.pth'),
config=SimpleNamespace(crop_ratio=0.55, input_size=256, num_modules=2,
hg_num_features=256, hg_depth=4, use_avg_pool=False))
hg_num_features=256, hg_depth=4, use_avg_pool=False,
use_instance_norm=False))
else:
raise ValueError('name must be set to either 2dfan2, 2dfan4, or 2dfan2_alt')

Expand Down

0 comments on commit fa410a8

Please sign in to comment.