原始网络vgg19(出来512x7x7)加三个分支的全连接网络。
修改 先采用resnet19替换vgg19
直接采用tochvision的resnet19网络和预训练模型,注意预训练网络是一个完整的resnet网络,网络最后两层包括一个池化层和一个全连接层,我们是不需要的所以要删去网络的最后两层。
import torchvision.models as models
resnet18 = models.resnet18(pretrained=True)
resnet18 = torch.nn.Sequential(*(list(resnet18.children())[:-2]))
model = Model(features=resnet18)
这样就将原来的特征提取由vgg变成了resnet
接下来是全连接层的替换,在这之前不太确定动全连接层会不会导致网络失效,所以先采取1x1的卷积来降低最后通道的数量
网络添加如下层:
self.conv1 = nn.Conv2d(512, 256, kernel_size=1, stride=1, padding=0)
self.bn1 = nn.BatchNorm2d(256, momentum=0.1)
self.conv2 = nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0)
self.bn2 = nn.BatchNorm2d(128, momentum=0.1)
self.relu = nn.ReLU(inplace=True)
x = self.features(x) # 512 x 7 x 7
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.conv2(x)
x = self.bn2(x)
x = self.relu(x) # 128 x 7 x 7
可以看出经过1x1的卷积,通道数由原来的512降低为128,不同通道之间的特征进行了融合。网络的大小也由最初的352m降低至现在的105m缩减了约1/3。
最后一步也就是三个全连接层的合并,可以看在本网络中网络的大小与全连接层的参数大小息息相关,而且预测的9个参数可以合并到一个全连接层来训练。
首先我们设计一个新的全连接层:
self.fc1 = nn.Sequential(
nn.Linear(128 * 7 * 7, 128),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(128, 128),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(128, 3+3*bins)
)
然后将resnet出来的参数直接送入fc1
x = x.view(-1, 128 * 7 * 7)
allPar = self.fc1(x)
dim = []
orient = []
conf = []
for ten in allPar:
len1 = len(ten) - 3
orient.append(ten[3: int(3 + 2 * len1 / 3)])
conf.append(ten[int(3 + 2 * len1 / 3): len1 + 3])
dim.append(ten[0:3])
dim = torch.stack(dim)
conf = torch.stack(conf)
orient = torch.stack(orient)
orient = orient.view(-1, self.bins, 2)
orient = F.normalize(orient, dim=2)
#print(dim.shape)
#print(conf.shape)
#print(orient.shape)
return orient, conf, dim
为了和后面loss的处理适配,我们还要依据原来网络的样子,拆分参数,如上所示。
最终得到修改后网络。