15090929758

Committed 19 Oct 2024 05:10AM UTC coverage: 54.553%. First build

Build # 15090929758

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #228 from cosanlab/huggingface

WIP: Huggingface Integration

Run Details

702 of 1620 new or added lines in 46 files covered. (43.33%)

3409 of 6249 relevant lines covered (54.55%)

3.27 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

86.43

/feat/landmark_detectors/mobilefacenet_test.py

from torch.nn import (
    Linear,
    Conv2d,
    BatchNorm1d,
    BatchNorm2d,
    PReLU,
    Sequential,
    Module,
)
import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin

##################################  Original Arcface Model #############################################################


class Flatten(Module):
    def forward(self, input):
        return input.view(input.size(0), -1)


##################################  MobileFaceNet #############################################################


class Conv_block(Module):
    def __init__(
        self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1
    ):
        super(Conv_block, self).__init__()
        self.conv = Conv2d(
            in_c,
            out_channels=out_c,
            kernel_size=kernel,
            groups=groups,
            stride=stride,
            padding=padding,
            bias=False,
        )
        self.bn = BatchNorm2d(
            out_c
        )  # Here is another MPS issue where data is not float32
        self.prelu = PReLU(out_c)

        # Ensure BatchNorm parameters are float32
        self.bn.weight.data = self.bn.weight.data.float()
        self.bn.bias.data = self.bn.bias.data.float()
        self.bn.running_mean.data = self.bn.running_mean.data.float()
        self.bn.running_var.data = self.bn.running_var.data.float()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.prelu(x)
        return x

    @property
    def weight(self):
        return self.conv.weight

    @property
    def bias(self):
        return self.conv.bias


class Linear_block(Module):
    def __init__(
        self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1
    ):
        super(Linear_block, self).__init__()
        self.conv = Conv2d(
            in_c,
            out_channels=out_c,
            kernel_size=kernel,
            groups=groups,
            stride=stride,
            padding=padding,
            bias=False,
        )
        self.bn = BatchNorm2d(out_c)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x


class Depth_Wise(Module):
    def __init__(
        self,
        in_c,
        out_c,
        residual=False,
        kernel=(3, 3),
        stride=(2, 2),
        padding=(1, 1),
        groups=1,
    ):
        super(Depth_Wise, self).__init__()
        self.conv = Conv_block(
            in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1)
        )
        self.conv_dw = Conv_block(
            groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride
        )
        self.project = Linear_block(
            groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1)
        )
        self.residual = residual

    def forward(self, x):
        if self.residual:
            short_cut = x
        x = self.conv(x)
        x = self.conv_dw(x)
        x = self.project(x)
        if self.residual:
            output = short_cut + x
        else:
            output = x
        return output


class Residual(Module):
    def __init__(
        self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)
    ):
        super(Residual, self).__init__()
        modules = []
        for _ in range(num_block):
            modules.append(
                Depth_Wise(
                    c,
                    c,
                    residual=True,
                    kernel=kernel,
                    padding=padding,
                    stride=stride,
                    groups=groups,
                )
            )
        self.model = Sequential(*modules)

    def forward(self, x):
        return self.model(x)


class GNAP(Module):
    def __init__(self, embedding_size):
        super(GNAP, self).__init__()
        assert embedding_size == 512
        self.bn1 = BatchNorm2d(512, affine=False)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))

        self.bn2 = BatchNorm1d(512, affine=False)

    def forward(self, x):
        x = self.bn1(x)
        x_norm = torch.norm(x, 2, 1, True)
        x_norm_mean = torch.mean(x_norm)
        weight = x_norm_mean / x_norm
        x = x * weight
        x = self.pool(x)
        x = x.view(x.shape[0], -1)
        feature = self.bn2(x)
        return feature


class GDC(Module):
    def __init__(self, embedding_size):
        super(GDC, self).__init__()
        self.conv_6_dw = Linear_block(
            512, 512, groups=512, kernel=(7, 7), stride=(1, 1), padding=(0, 0)
        )
        self.conv_6_flatten = Flatten()
        self.linear = Linear(512, embedding_size, bias=False)
        # self.bn = BatchNorm1d(embedding_size, affine=False)
        self.bn = BatchNorm1d(embedding_size)

    def forward(self, x):
        x = self.conv_6_dw(x)
        x = self.conv_6_flatten(x)
        x = self.linear(x)
        x = self.bn(x)
        return x


class MobileFaceNet(Module, PyTorchModelHubMixin):
    def __init__(self, input_size, embedding_size=512, output_name="GDC", device="cpu"):
        super(MobileFaceNet, self).__init__()
        # Make sure this module is compatible with mps
        self.device = device
        self.to(device)
        # torch.set_default_dtype(torch.float32) # Ensure default dtype is float32 for MPS compatibility

        assert output_name in ["GNAP", "GDC"]
        assert input_size[0] in [112]
        self.conv1 = Conv_block(3, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1))
        self.conv2_dw = Conv_block(
            64, 64, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64
        )
        self.conv_23 = Depth_Wise(
            64, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128
        )
        self.conv_3 = Residual(
            64, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1)
        )
        self.conv_34 = Depth_Wise(
            64, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256
        )
        self.conv_4 = Residual(
            128, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)
        )
        self.conv_45 = Depth_Wise(
            128, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512
        )
        self.conv_5 = Residual(
            128, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)
        )
        self.conv_6_sep = Conv_block(
            128, 512, kernel=(1, 1), stride=(1, 1), padding=(0, 0)
        )
        if output_name == "GNAP":
            self.output_layer = GNAP(512)
        else:
            self.output_layer = GDC(embedding_size)

        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                if m.bias is not None:
                    m.bias.data.zero_()

    def forward(self, x):
        # Ensure this module is compatible with mps
        x = x.to(self.device)
        x = x.to(self.device).float()

        out = self.conv1(x)

        out = self.conv2_dw(out)

        out = self.conv_23(out)

        out = self.conv_3(out)

        out = self.conv_34(out)

        out = self.conv_4(out)

        out = self.conv_45(out)

        out = self.conv_5(out)

        conv_features = self.conv_6_sep(out)
        out = self.output_layer(conv_features)
        return out, conv_features

1	from torch.nn import (	6✔
2	Linear,
3	Conv2d,
4	BatchNorm1d,
5	BatchNorm2d,
6	PReLU,
7	Sequential,
8	Module,
9	)
10	import torch	6✔
11	import torch.nn as nn	6✔
12	from huggingface_hub import PyTorchModelHubMixin	6✔
13
14	################################## Original Arcface Model #############################################################
15
16
17	class Flatten(Module):	6✔
18	def forward(self, input):	6✔
19	return input.view(input.size(0), -1)	6✔
20
21
22	################################## MobileFaceNet #############################################################
23
24
25	class Conv_block(Module):	6✔
26	def __init__(	6✔
27	self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1
28	):
29	super(Conv_block, self).__init__()	6✔
30	self.conv = Conv2d(	6✔
31	in_c,
32	out_channels=out_c,
33	kernel_size=kernel,
34	groups=groups,
35	stride=stride,
36	padding=padding,
37	bias=False,
38	)
39	self.bn = BatchNorm2d(	6✔
40	out_c
41	) # Here is another MPS issue where data is not float32
42	self.prelu = PReLU(out_c)	6✔
43
44	# Ensure BatchNorm parameters are float32
45	self.bn.weight.data = self.bn.weight.data.float()	6✔
46	self.bn.bias.data = self.bn.bias.data.float()	6✔
47	self.bn.running_mean.data = self.bn.running_mean.data.float()	6✔
48	self.bn.running_var.data = self.bn.running_var.data.float()	6✔
49
50	def forward(self, x):	6✔
51	x = self.conv(x)	6✔
52	x = self.bn(x)	6✔
53	x = self.prelu(x)	6✔
54	return x	6✔
55
56	@property	6✔
57	def weight(self):	6✔
NEW 58	return self.conv.weight	×
59
60	@property	6✔
61	def bias(self):	6✔
NEW 62	return self.conv.bias	×
63
64
65	class Linear_block(Module):	6✔
66	def __init__(	6✔
67	self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1
68	):
69	super(Linear_block, self).__init__()	6✔
70	self.conv = Conv2d(	6✔
71	in_c,
72	out_channels=out_c,
73	kernel_size=kernel,
74	groups=groups,
75	stride=stride,
76	padding=padding,
77	bias=False,
78	)
79	self.bn = BatchNorm2d(out_c)	6✔
80
81	def forward(self, x):	6✔
82	x = self.conv(x)	6✔
83	x = self.bn(x)	6✔
84	return x	6✔
85
86
87	class Depth_Wise(Module):	6✔
88	def __init__(	6✔
89	self,
90	in_c,
91	out_c,
92	residual=False,
93	kernel=(3, 3),
94	stride=(2, 2),
95	padding=(1, 1),
96	groups=1,
97	):
98	super(Depth_Wise, self).__init__()	6✔
99	self.conv = Conv_block(	6✔
100	in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1)
101	)
102	self.conv_dw = Conv_block(	6✔
103	groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride
104	)
105	self.project = Linear_block(	6✔
106	groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1)
107	)
108	self.residual = residual	6✔
109
110	def forward(self, x):	6✔
111	if self.residual:	6✔
112	short_cut = x	6✔
113	x = self.conv(x)	6✔
114	x = self.conv_dw(x)	6✔
115	x = self.project(x)	6✔
116	if self.residual:	6✔
117	output = short_cut + x	6✔
118	else:
119	output = x	6✔
120	return output	6✔
121
122
123	class Residual(Module):	6✔
124	def __init__(	6✔
125	self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)
126	):
127	super(Residual, self).__init__()	6✔
128	modules = []	6✔
129	for _ in range(num_block):	6✔
130	modules.append(	6✔
131	Depth_Wise(
132	c,
133	c,
134	residual=True,
135	kernel=kernel,
136	padding=padding,
137	stride=stride,
138	groups=groups,
139	)
140	)
141	self.model = Sequential(*modules)	6✔
142
143	def forward(self, x):	6✔
144	return self.model(x)	6✔
145
146
147	class GNAP(Module):	6✔
148	def __init__(self, embedding_size):	6✔
149	super(GNAP, self).__init__()	×
150	assert embedding_size == 512	×
151	self.bn1 = BatchNorm2d(512, affine=False)	×
152	self.pool = nn.AdaptiveAvgPool2d((1, 1))	×
153
154	self.bn2 = BatchNorm1d(512, affine=False)	×
155
156	def forward(self, x):	6✔
157	x = self.bn1(x)	×
158	x_norm = torch.norm(x, 2, 1, True)	×
159	x_norm_mean = torch.mean(x_norm)	×
160	weight = x_norm_mean / x_norm	×
161	x = x * weight	×
162	x = self.pool(x)	×
163	x = x.view(x.shape[0], -1)	×
164	feature = self.bn2(x)	×
165	return feature	×
166
167
168	class GDC(Module):	6✔
169	def __init__(self, embedding_size):	6✔
170	super(GDC, self).__init__()	6✔
171	self.conv_6_dw = Linear_block(	6✔
172	512, 512, groups=512, kernel=(7, 7), stride=(1, 1), padding=(0, 0)
173	)
174	self.conv_6_flatten = Flatten()	6✔
175	self.linear = Linear(512, embedding_size, bias=False)	6✔
176	# self.bn = BatchNorm1d(embedding_size, affine=False)
177	self.bn = BatchNorm1d(embedding_size)	6✔
178
179	def forward(self, x):	6✔
180	x = self.conv_6_dw(x)	6✔
181	x = self.conv_6_flatten(x)	6✔
182	x = self.linear(x)	6✔
183	x = self.bn(x)	6✔
184	return x	6✔
185
186
187	class MobileFaceNet(Module, PyTorchModelHubMixin):	6✔
188	def __init__(self, input_size, embedding_size=512, output_name="GDC", device="cpu"):	6✔
189	super(MobileFaceNet, self).__init__()	6✔
190	# Make sure this module is compatible with mps
191	self.device = device	6✔
192	self.to(device)	6✔
193	# torch.set_default_dtype(torch.float32) # Ensure default dtype is float32 for MPS compatibility
194
195	assert output_name in ["GNAP", "GDC"]	6✔
196	assert input_size[0] in [112]	6✔
197	self.conv1 = Conv_block(3, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1))	6✔
198	self.conv2_dw = Conv_block(	6✔
199	64, 64, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64
200	)
201	self.conv_23 = Depth_Wise(	6✔
202	64, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128
203	)
204	self.conv_3 = Residual(	6✔
205	64, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1)
206	)
207	self.conv_34 = Depth_Wise(	6✔
208	64, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256
209	)
210	self.conv_4 = Residual(	6✔
211	128, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)
212	)
213	self.conv_45 = Depth_Wise(	6✔
214	128, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512
215	)
216	self.conv_5 = Residual(	6✔
217	128, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)
218	)
219	self.conv_6_sep = Conv_block(	6✔
220	128, 512, kernel=(1, 1), stride=(1, 1), padding=(0, 0)
221	)
222	if output_name == "GNAP":	6✔
223	self.output_layer = GNAP(512)	×
224	else:
225	self.output_layer = GDC(embedding_size)	6✔
226
227	self._initialize_weights()	6✔
228
229	def _initialize_weights(self):	6✔
230	for m in self.modules():	6✔
231	if isinstance(m, nn.Conv2d):	6✔
232	nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")	6✔
233	if m.bias is not None:	6✔
234	m.bias.data.zero_()	×
235	elif isinstance(m, nn.BatchNorm2d):	6✔
236	m.weight.data.fill_(1)	6✔
237	m.bias.data.zero_()	6✔
238	elif isinstance(m, nn.Linear):	6✔
239	nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")	6✔
240	if m.bias is not None:	6✔
241	m.bias.data.zero_()	×
242
243	def forward(self, x):	6✔
244	# Ensure this module is compatible with mps
245	x = x.to(self.device)	6✔
246	x = x.to(self.device).float()	6✔
247
248	out = self.conv1(x)	6✔
249
250	out = self.conv2_dw(out)	6✔
251
252	out = self.conv_23(out)	6✔
253
254	out = self.conv_3(out)	6✔
255
256	out = self.conv_34(out)	6✔
257
258	out = self.conv_4(out)	6✔
259
260	out = self.conv_45(out)	6✔
261
262	out = self.conv_5(out)	6✔
263
264	conv_features = self.conv_6_sep(out)	6✔
265	out = self.output_layer(conv_features)	6✔
266	return out, conv_features	6✔

cosanlab / py-feat / 15090929758

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous