-
Notifications
You must be signed in to change notification settings - Fork 53
/
espnetv2.patch
92 lines (85 loc) · 4.8 KB
/
espnetv2.patch
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
diff -x .git -x .gitignore -x '*.png' -x '*.pth' -x '*.idea*' -r -N -r -u ./segmentation/cnn/SegmentationModel.py esanet/src/models/external_code/ESPNetv2/segmentation/cnn/SegmentationModel.py
--- ./segmentation/cnn/SegmentationModel.py 2020-12-06 10:15:23.000000000 +0100
+++ esanet/src/models/external_code/ESPNetv2/segmentation/cnn/SegmentationModel.py 2020-12-04 17:56:52.000000000 +0100
@@ -47,25 +47,30 @@
def hierarchicalUpsample(self, x, factor=3):
for i in range(factor):
- x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
+ size = (2 * int(x.shape[2]), 2 * int(x.shape[3]))
+ x = F.interpolate(x, size=size, mode='bilinear', align_corners=False)
return x
def forward(self, input):
out_l1, out_l2, out_l3, out_l4 = self.net(input, seg=True)
out_l4_proj = self.proj_L4_C(out_l4)
- up_l4_to_l3 = F.interpolate(out_l4_proj, scale_factor=2, mode='bilinear', align_corners=True)
+ size = (2*int(out_l4_proj.shape[2]), 2*int(out_l4_proj.shape[3]))
+ up_l4_to_l3 = F.interpolate(out_l4_proj, size=size, mode='bilinear', align_corners=False)
merged_l3_upl4 = self.pspMod(torch.cat([out_l3, up_l4_to_l3], 1))
proj_merge_l3_bef_act = self.project_l3(merged_l3_upl4)
proj_merge_l3 = self.act_l3(proj_merge_l3_bef_act)
- out_up_l3 = F.interpolate(proj_merge_l3, scale_factor=2, mode='bilinear', align_corners=True)
+ size = (2*int(proj_merge_l3.shape[2]), 2*int(proj_merge_l3.shape[3]))
+ out_up_l3 = F.interpolate(proj_merge_l3, size=size, mode='bilinear', align_corners=False)
merge_l2 = self.project_l2(torch.cat([out_l2, out_up_l3], 1))
- out_up_l2 = F.interpolate(merge_l2, scale_factor=2, mode='bilinear', align_corners=True)
+ size = (2 * int(merge_l2.shape[2]), 2 * int(merge_l2.shape[3]))
+ out_up_l2 = F.interpolate(merge_l2, size=size, mode='bilinear', align_corners=False)
merge_l1 = self.project_l1(torch.cat([out_l1, out_up_l2], 1))
if self.training:
- return F.interpolate(merge_l1, scale_factor=2, mode='bilinear', align_corners=True), self.hierarchicalUpsample(proj_merge_l3_bef_act)
+ return F.interpolate(merge_l1, scale_factor=2, mode='bilinear', align_corners=False), self.hierarchicalUpsample(proj_merge_l3_bef_act)
else:
- return F.interpolate(merge_l1, scale_factor=2, mode='bilinear', align_corners=True)
+ size = (2 * int(merge_l1.shape[2]), 2 * int(merge_l1.shape[3]))
+ return F.interpolate(merge_l1, size=size, mode='bilinear', align_corners=False)
if __name__ == '__main__':
diff -x .git -x .gitignore -x '*.png' -x '*.pth' -x '*.idea*' -r -N -r -u ./segmentation/cnn/cnn_utils.py esanet/src/models/external_code/ESPNetv2/segmentation/cnn/cnn_utils.py
--- ./segmentation/cnn/cnn_utils.py 2020-12-06 10:15:23.000000000 +0100
+++ esanet/src/models/external_code/ESPNetv2/segmentation/cnn/cnn_utils.py 2020-12-04 17:56:52.000000000 +0100
@@ -16,11 +16,11 @@
self.project = CBR(features * (len(sizes) + 1), out_features, 1, 1)
def forward(self, feats):
- h, w = feats.size(2), feats.size(3)
+ h, w = int(feats.size(2)), int(feats.size(3))
out = [feats]
for stage in self.stages:
feats = F.avg_pool2d(feats, kernel_size=3, stride=2, padding=1)
- upsampled = F.interpolate(input=stage(feats), size=(h, w), mode='bilinear', align_corners=True)
+ upsampled = F.interpolate(input=stage(feats), size=(h, w), mode='bilinear', align_corners=False)
out.append(upsampled)
return self.project(torch.cat(out, dim=1))
diff -x .git -x .gitignore -x '*.png' -x '*.pth' -x '*.idea*' -r -N -r -u ./segmentation/to_onnx.py esanet/src/models/external_code/ESPNetv2/segmentation/to_onnx.py
--- ./segmentation/to_onnx.py 1970-01-01 01:00:00.000000000 +0100
+++ esanet/src/models/external_code/ESPNetv2/segmentation/to_onnx.py 2020-12-04 17:56:52.000000000 +0100
@@ -0,0 +1,31 @@
+import os
+import torch
+from cnn.SegmentationModel import EESPNet_Seg
+
+N_CLASSES = 19
+H = 512
+W = 1024
+
+model = EESPNet_Seg(classes=N_CLASSES)
+model.eval()
+
+rgb = torch.rand(size=(1, 3, H, W), dtype=torch.float32)
+
+device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+model.to(device)
+rgb = rgb.to(device)
+
+out_dir = '../../onnx_models'
+os.makedirs(out_dir, exist_ok=True)
+onnx_file_path = os.path.join(out_dir, 'espnetv2.onnx')
+
+torch.onnx.export(model,
+ rgb,
+ onnx_file_path,
+ export_params=True,
+ input_names=['rgb'],
+ output_names=['output'],
+ do_constant_folding=True,
+ verbose=False,
+ opset_version=10
+ )