Spaces:
Runtime error
Runtime error
File size: 3,271 Bytes
2e34814 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
from torch.nn import (
Linear,
Conv2d,
BatchNorm1d,
BatchNorm2d,
PReLU,
Dropout,
Sequential,
Module,
)
from criteria.helpers import (
get_blocks,
Flatten,
bottleneck_IR,
bottleneck_IR_SE,
l2_norm,
)
"""
Modified Backbone implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch)
"""
class Backbone(Module):
def __init__(self, input_size, num_layers, mode="ir", drop_ratio=0.4, affine=True):
super(Backbone, self).__init__()
assert input_size in [112, 224], "input_size should be 112 or 224"
assert num_layers in [50, 100, 152], "num_layers should be 50, 100 or 152"
assert mode in ["ir", "ir_se"], "mode should be ir or ir_se"
blocks = get_blocks(num_layers)
if mode == "ir":
unit_module = bottleneck_IR
elif mode == "ir_se":
unit_module = bottleneck_IR_SE
self.input_layer = Sequential(
Conv2d(3, 64, (3, 3), 1, 1, bias=False), BatchNorm2d(64), PReLU(64)
)
if input_size == 112:
self.output_layer = Sequential(
BatchNorm2d(512),
Dropout(drop_ratio),
Flatten(),
Linear(512 * 7 * 7, 512),
BatchNorm1d(512, affine=affine),
)
else:
self.output_layer = Sequential(
BatchNorm2d(512),
Dropout(drop_ratio),
Flatten(),
Linear(512 * 14 * 14, 512),
BatchNorm1d(512, affine=affine),
)
modules = []
for block in blocks:
for bottleneck in block:
modules.append(
unit_module(
bottleneck.in_channel, bottleneck.depth, bottleneck.stride
)
)
self.body = Sequential(*modules)
def forward(self, x):
x = self.input_layer(x)
x = self.body(x)
x = self.output_layer(x)
return l2_norm(x)
def IR_50(input_size):
"""Constructs a ir-50 model."""
model = Backbone(input_size, num_layers=50, mode="ir", drop_ratio=0.4, affine=False)
return model
def IR_101(input_size):
"""Constructs a ir-101 model."""
model = Backbone(
input_size, num_layers=100, mode="ir", drop_ratio=0.4, affine=False
)
return model
def IR_152(input_size):
"""Constructs a ir-152 model."""
model = Backbone(
input_size, num_layers=152, mode="ir", drop_ratio=0.4, affine=False
)
return model
def IR_SE_50(input_size):
"""Constructs a ir_se-50 model."""
model = Backbone(
input_size, num_layers=50, mode="ir_se", drop_ratio=0.4, affine=False
)
return model
def IR_SE_101(input_size):
"""Constructs a ir_se-101 model."""
model = Backbone(
input_size, num_layers=100, mode="ir_se", drop_ratio=0.4, affine=False
)
return model
def IR_SE_152(input_size):
"""Constructs a ir_se-152 model."""
model = Backbone(
input_size, num_layers=152, mode="ir_se", drop_ratio=0.4, affine=False
)
return model
|