-
Notifications
You must be signed in to change notification settings - Fork 0
/
feature_distill_simple.py
102 lines (93 loc) · 4.03 KB
/
feature_distill_simple.py
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
# --------------------------------------------------------
# Tensorflow Faster R-CNN
# Licensed under The MIT License [see LICENSE for details]
# Written by Xinlei Chen
# --------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math
import torchvision.models as models
import pdb
def conv3x3(in_planes, out_planes, pad, dilation, stride=1):
"3x3 convolution with padding"
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=pad, dilation=dilation, bias=True)
def conv1x1(in_planes, out_planes, stride=1):
"3x3 convolution with padding"
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,
padding=0, bias=False)
class F_D(nn.Module):
def __init__(self, input_channel, output_channel=3):
super(F_D, self).__init__()
self.reduced_conv = conv3x3(input_channel, 64, 1, 1)
self.reduced_bn = nn.BatchNorm2d(64)
self.ddc_x1 = conv3x3(64, 64, 2, 2)
self.ddc_bn1 = nn.BatchNorm2d(64)
self.ddc_x2 = conv3x3(128, 64, 4, 4)
self.ddc_bn2 = nn.BatchNorm2d(64)
self.ddc_x3 = conv3x3(192, 64, 8, 8)
self.ddc_bn3 = nn.BatchNorm2d(64)
self.post_conv = conv1x1(256, 128)
self.post_bn = nn.BatchNorm2d(128)
self.pool1_conv = conv1x1(128, 64)
self.pool1_bn = nn.BatchNorm2d(64)
self.pool2_conv = conv1x1(128, 64)
self.pool2_bn = nn.BatchNorm2d(64)
self.pool3_conv = conv1x1(128, 64)
self.pool3_bn = nn.BatchNorm2d(64)
self.pool4_conv = conv1x1(128, 64)
self.pool4_bn = nn.BatchNorm2d(64)
self.conv2 = conv1x1(256, 128)
self.bn2 = nn.BatchNorm2d(128)
self.conv_cls = conv1x1(128, output_channel)
#self.fc = nn.Linear(128, 2)
def forward(self, x):
# reduced_x
x = F.interpolate(x, size=(256, 256), mode='bilinear')
x_r = F.relu(self.reduced_bn(self.reduced_conv(x)))
#ddc x1
x1 = F.relu(self.ddc_bn1(self.ddc_x1(x_r)))
x1_c = torch.cat((x_r, x1), 1)
# ddc x2
x2 = F.relu(self.ddc_bn2(self.ddc_x2(x1_c)))
x2_c = torch.cat((x1_c, x2), 1)
# ddc x3
x3 = F.relu(self.ddc_bn3(self.ddc_x3(x2_c)))
#all concat
x1_p = torch.cat((x_r, x1), 1)
x2_p = torch.cat((x1_p, x2), 1)
x3_p = torch.cat((x2_p, x3), 1) #[1,256,256,256]
#post layers
x_post = F.relu(self.post_bn(self.post_conv(x3_p)))
# First level
x_b_1 = F.avg_pool2d(x_post, (x_post.size(2) // 64, x_post.size(3) // 64))
x_b_1 = F.relu(self.pool3_bn(self.pool3_conv(x_b_1))) #[b,64,64,64]
# Second level
x_b_2 = F.avg_pool2d(x_post, (x_post.size(2) // 128, x_post.size(3) // 128))
x_b_2 = F.relu(self.pool4_bn(self.pool4_conv(x_b_2))) #[b,64,128,128]
#unsampling layer
x_b_1_u = F.interpolate(x_b_1, size=(x_post.size(2), x_post.size(3)), mode='bilinear') #[4,64,224,224]
x_b_2_u = F.interpolate(x_b_2, size=(x_post.size(2), x_post.size(3)), mode='bilinear') #[4,64,224,224]
#concat layer
x_c_1 = torch.cat((x_post,x_b_1_u),1)
x_c_2 = torch.cat((x_c_1, x_b_2_u), 1)
# x_c_3 = torch.cat((x_c_2, x_b_2_u), 1)
# x_c_4 = torch.cat((x_c_3, x_b_1_u), 1) #[4,384,224,224]
#domain classifier
x_p = F.relu(self.bn2(self.conv2(x_c_2)))
#x = F.avg_pool2d(x_p, (x_p.size(2), x_p.size(3)))
#x = x.view(-1, 128)
x = self.conv_cls(x_p)
return x
if __name__ == "__main__":
a=torch.randn([4,18,224,224])
dis_model = F_D(input_channel=18)
x=dis_model(a)
#/remote-home/share/42/cyc19307140030/yolov5/feature_distill/
#dis_model.save('/remote-home/share/42/cyc19307140030/yolov5/feature_distill/ddcpp_F_D.pth')
torch.save(dis_model,'/remote-home/share/42/cyc19307140030/yolov5/feature_distill/ddcpp_F_D.pth')