Skip to content

Commit 19b7ca8

Browse files
chyomin06fracape
chyomin06
authored and
fracape
committed
[refactor] invconv commit that supports jde yolo with integer convolutional layers
1 parent 8345f11 commit 19b7ca8

File tree

4 files changed

+317
-1
lines changed

4 files changed

+317
-1
lines changed

‎cfgs/vision_model/default.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ jde_1088x608:
3636
model_path_prefix: ${..model_root_path}
3737
cfg: "models/Towards-Realtime-MOT/cfg/yolov3_1088x608.cfg"
3838
weights: "weights/jde/jde.1088x608.uncertainty.pt"
39+
integer_conv_weight: False
3940
iou_thres: 0.5
4041
conf_thres: 0.5
4142
nms_thres: 0.4

‎compressai_vision/model_wrappers/intconv_wrapper.py‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
# Copyright (c) 2022-2024, InterDigital Communications, Inc
2+
# All rights reserved.
3+
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted (subject to the limitations in the disclaimer
6+
# below) provided that the following conditions are met:
7+
8+
# * Redistributions of source code must retain the above copyright notice,
9+
# this list of conditions and the following disclaimer.
10+
# * Redistributions in binary form must reproduce the above copyright notice,
11+
# this list of conditions and the following disclaimer in the documentation
12+
# and/or other materials provided with the distribution.
13+
# * Neither the name of InterDigital Communications, Inc nor the names of its
14+
# contributors may be used to endorse or promote products derived from this
15+
# software without specific prior written permission.
16+
17+
# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
18+
# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
19+
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT
20+
# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
21+
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
22+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
23+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
24+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
25+
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
26+
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
27+
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
28+
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
30+
import logging
31+
32+
import numpy as np
33+
import torch
34+
from torch import nn
35+
36+
37+
class IntConv2dWrapper(nn.Conv2d):
38+
def __init__(
39+
self,
40+
in_channels: int,
41+
out_channels: int,
42+
kernel_size,
43+
stride=1,
44+
padding=0,
45+
dilation=1,
46+
groups: int = 1,
47+
bias: bool = True,
48+
padding_mode: str = "zeros",
49+
device=None,
50+
dtype=None,
51+
) -> None:
52+
super().__init__(
53+
in_channels,
54+
out_channels,
55+
kernel_size,
56+
stride,
57+
padding,
58+
dilation,
59+
groups,
60+
bias,
61+
padding_mode,
62+
device,
63+
dtype,
64+
)
65+
self.initified_weight_mode = False
66+
67+
"""
68+
def _set_mode(mode):
69+
global _precision, _high_precision, _mode
70+
71+
if mode == 'none':
72+
_precision = 0
73+
elif mode == 'float32':
74+
_precision = 2**(23+1)
75+
elif mode == 'float64':
76+
_precision = 2**(52+1)
77+
78+
_mode = mode
79+
torch.backends.cudnn.enabled = mode=='none'
80+
"""
81+
82+
def quantize_weights(self):
83+
self.initified_weight_mode = True
84+
85+
if self.bias is None:
86+
self.float_bias = torch.zeros(self.out_channels, device=self.weight.device)
87+
else:
88+
self.float_bias = self.bias.detach().clone()
89+
90+
if self.weight.dtype == torch.float32:
91+
_precision = 2 ** (23 + 1)
92+
elif self.weight.dtype == torch.float64:
93+
_precision = 2 ** (52 + 1)
94+
else:
95+
logging.warning(
96+
f"Unsupported dtype {self.weight.dtype}. Behaviour may lead unexpected results."
97+
)
98+
_precision = 2 ** (23 + 1)
99+
100+
###### REFERENCE FROM VCMRMS ######
101+
# sf const
102+
sf_const = 48
103+
104+
N = np.prod(self.weight.shape[1:])
105+
self.N = N
106+
self.factor = np.sqrt(_precision)
107+
# self.sf = 1/6 #precision bits allocation factor
108+
self.sf = np.sqrt(sf_const / N)
109+
110+
# perform the calculate ion CPU to stabalize the calculation
111+
self.w_sum = self.weight.cpu().abs().sum(axis=[1, 2, 3]).to(self.weight.device)
112+
self.w_sum[self.w_sum == 0] = 1 # prevent divide by 0
113+
114+
self.fw = (self.factor / self.sf - np.sqrt(N / 12) * 5) / self.w_sum
115+
116+
# intify weights
117+
self.weight.requires_grad = False # Just make sure
118+
self.weight.copy_(
119+
torch.round(self.weight.detach().clone() * self.fw.view(-1, 1, 1, 1))
120+
)
121+
122+
# set bias to 0
123+
if self.bias is not None:
124+
self.bias.requires_grad = False # Just make sure
125+
self.bias.zero_()
126+
127+
###### END OF REFERENCE FROM VCMRMS ######
128+
129+
def forward(self, x: torch.Tensor):
130+
if not self.initified_weight_mode:
131+
return super().forward(x)
132+
133+
_dtype = x.dtype
134+
_cudnn_enabled = torch.backends.cudnn.enabled
135+
torch.backends.cudnn.enabled = False
136+
137+
###### REFERENCE FROM VCMRMS ######
138+
139+
# Calculate factor
140+
fx = 1
141+
142+
x_abs = x.abs()
143+
x_max = x_abs.max()
144+
if x_max > 0:
145+
fx = (self.factor * self.sf - 0.5) / x_max
146+
147+
# intify x
148+
x = torch.round(fx * x)
149+
x = super().forward(x)
150+
151+
# x should be all integers
152+
x /= fx * self.fw.view(-1, 1, 1)
153+
x = x.float()
154+
155+
# apply bias in float format
156+
x = (x.permute(0, 2, 3, 1) + self.float_bias).permute(0, 3, 1, 2).contiguous()
157+
###### REFERENCE FROM VCMRMS ######
158+
159+
torch.backends.cudnn.enabled = _cudnn_enabled
160+
161+
return x.to(_dtype)

‎compressai_vision/model_wrappers/jde.py‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,12 +43,16 @@
4343
)
4444
from jde.utils.kalman_filter import KalmanFilter
4545
from jde.utils.utils import non_max_suppression, scale_coords
46-
from torch import Tensor
4746

4847
from compressai_vision.registry import register_vision_model
4948

5049
from .base_wrapper import BaseWrapper
5150

51+
# Patch in modified create_modules
52+
from .jde_lowlevel import create_modules
53+
54+
jde.models.create_modules = create_modules
55+
5256
__all__ = [
5357
"jde_1088x608",
5458
]
@@ -84,6 +88,8 @@ def __init__(self, device: str, **kwargs):
8488
self.model_configs["frame_rate"] / 30.0 * self.model_configs["track_buffer"]
8589
)
8690

91+
integer_conv_weight = bool(kwargs["integer_conv_weight"])
92+
8793
assert "splits" in kwargs, "Split layer ids must be provided"
8894
self.split_layer_list = kwargs["splits"]
8995
self.features_at_splits = dict(
@@ -96,6 +102,12 @@ def __init__(self, device: str, **kwargs):
96102
strict=False,
97103
)
98104
self.darknet.to(device).eval()
105+
for param in self.darknet.parameters():
106+
param.requires_grad = False
107+
108+
# must be called after loading weights to a model
109+
if integer_conv_weight:
110+
self.darknet = self.quantize_weights(self.darknet)
99111

100112
self.kalman_filter = KalmanFilter()
101113

@@ -114,6 +126,17 @@ def reset(self):
114126

115127
self.frame_id = 0
116128

129+
@staticmethod
130+
def quantize_weights(model):
131+
132+
for module_def, module in zip(model.module_defs, model.module_list):
133+
if module_def["type"] == "convolutional":
134+
for m in module:
135+
if type(m).__name__ == "IntConv2dWrapper":
136+
m.quantize_weights()
137+
138+
return model
139+
117140
def input_to_features(self, x, device: str) -> Dict:
118141
"""Computes deep features at the intermediate layer(s) all the way from the input"""
119142
self.darknet = self.darknet.to(device).eval()

‎compressai_vision/model_wrappers/jde_lowlevel.py‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
# Copyright (c) 2022-2023, InterDigital Communications, Inc
2+
# All rights reserved.
3+
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted (subject to the limitations in the disclaimer
6+
# below) provided that the following conditions are met:
7+
8+
# * Redistributions of source code must retain the above copyright notice,
9+
# this list of conditions and the following disclaimer.
10+
# * Redistributions in binary form must reproduce the above copyright notice,
11+
# this list of conditions and the following disclaimer in the documentation
12+
# and/or other materials provided with the distribution.
13+
# * Neither the name of InterDigital Communications, Inc nor the names of its
14+
# contributors may be used to endorse or promote products derived from this
15+
# software without specific prior written permission.
16+
17+
# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
18+
# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
19+
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT
20+
# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
21+
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
22+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
23+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
24+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
25+
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
26+
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
27+
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
28+
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
30+
import torch.nn as nn
31+
from jde.models import EmptyLayer, Upsample, YOLOLayer
32+
33+
from .intconv_wrapper import IntConv2dWrapper
34+
35+
try:
36+
from jde.utils.syncbn import SyncBN
37+
38+
batch_norm = SyncBN # nn.BatchNorm2d
39+
except ImportError:
40+
batch_norm = nn.BatchNorm2d
41+
42+
43+
def create_modules(module_defs, device: str):
44+
"""
45+
Constructs module list of layer blocks from module configuration in module_defs
46+
"""
47+
hyperparams = module_defs.pop(0)
48+
output_filters = [int(hyperparams["channels"])]
49+
module_list = nn.ModuleList()
50+
yolo_layer_count = 0
51+
for i, module_def in enumerate(module_defs):
52+
modules = nn.Sequential()
53+
54+
if module_def["type"] == "convolutional":
55+
bn = int(module_def["batch_normalize"])
56+
filters = int(module_def["filters"])
57+
kernel_size = int(module_def["size"])
58+
pad = (kernel_size - 1) // 2 if int(module_def["pad"]) else 0
59+
modules.add_module(
60+
"conv_%d" % i,
61+
IntConv2dWrapper(
62+
in_channels=output_filters[-1],
63+
out_channels=filters,
64+
kernel_size=kernel_size,
65+
stride=int(module_def["stride"]),
66+
padding=pad,
67+
bias=not bn,
68+
),
69+
)
70+
if bn:
71+
after_bn = batch_norm(filters)
72+
modules.add_module("batch_norm_%d" % i, after_bn)
73+
# BN is uniformly initialized by default in pytorch 1.0.1.
74+
# In pytorch>1.2.0, BN weights are initialized with constant 1,
75+
# but we find with the uniform initialization the model converges faster.
76+
nn.init.uniform_(after_bn.weight)
77+
nn.init.zeros_(after_bn.bias)
78+
if module_def["activation"] == "leaky":
79+
modules.add_module("leaky_%d" % i, nn.LeakyReLU(0.1))
80+
81+
elif module_def["type"] == "maxpool":
82+
kernel_size = int(module_def["size"])
83+
stride = int(module_def["stride"])
84+
if kernel_size == 2 and stride == 1:
85+
modules.add_module("_debug_padding_%d" % i, nn.ZeroPad2d((0, 1, 0, 1)))
86+
maxpool = nn.MaxPool2d(
87+
kernel_size=kernel_size,
88+
stride=stride,
89+
padding=int((kernel_size - 1) // 2),
90+
)
91+
modules.add_module("maxpool_%d" % i, maxpool)
92+
93+
elif module_def["type"] == "upsample":
94+
upsample = Upsample(scale_factor=int(module_def["stride"]))
95+
modules.add_module("upsample_%d" % i, upsample)
96+
97+
elif module_def["type"] == "route":
98+
layers = [int(x) for x in module_def["layers"].split(",")]
99+
filters = sum([output_filters[i + 1 if i > 0 else i] for i in layers])
100+
modules.add_module("route_%d" % i, EmptyLayer())
101+
102+
elif module_def["type"] == "shortcut":
103+
filters = output_filters[int(module_def["from"])]
104+
modules.add_module("shortcut_%d" % i, EmptyLayer())
105+
106+
elif module_def["type"] == "yolo":
107+
anchor_idxs = [int(x) for x in module_def["mask"].split(",")]
108+
# Extract anchors
109+
anchors = [float(x) for x in module_def["anchors"].split(",")]
110+
anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
111+
anchors = [anchors[i] for i in anchor_idxs]
112+
nC = int(module_def["classes"]) # number of classes
113+
img_size = (int(hyperparams["width"]), int(hyperparams["height"]))
114+
# Define detection layer
115+
yolo_layer = YOLOLayer(
116+
anchors,
117+
nC,
118+
int(hyperparams["nID"]),
119+
int(hyperparams["embedding_dim"]),
120+
img_size,
121+
yolo_layer_count,
122+
device,
123+
)
124+
modules.add_module("yolo_%d" % i, yolo_layer)
125+
yolo_layer_count += 1
126+
127+
# Register module list and number of output filters
128+
module_list.append(modules)
129+
output_filters.append(filters)
130+
131+
return hyperparams, module_list

0 commit comments

Comments
 (0)