529 lines
23 KiB
Python
529 lines
23 KiB
Python
'''
|
||
author ： Hatcher W
|
||
Email： hcwang@semi.ac.cn
|
||
use: 密度估计网络结构
|
||
|
||
'''
|
||
import torch.nn as nn
|
||
from torchvision import models
|
||
from configs.configs import model_head
|
||
import torch
|
||
import os
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
os.environ['CUDA_VISIBLE_DEVICE'] = '6'
|
||
|
||
class MCNN(nn.Module):
|
||
def __init__(self, load_vgg=True):
|
||
super(MCNN, self).__init__()
|
||
#block1
|
||
self.layer1 = nn.Sequential(
|
||
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=9, padding=4),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(16),
|
||
nn.MaxPool2d((2, 2))
|
||
)
|
||
self.layer2 = nn.Sequential(
|
||
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=7, padding=3),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(32),
|
||
nn.MaxPool2d((2, 2))
|
||
)
|
||
self.layer3 = nn.Sequential(
|
||
nn.Conv2d(in_channels=32, out_channels=16, kernel_size=7, padding=3),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(16),
|
||
)
|
||
self.layer4 = nn.Sequential(
|
||
nn.Conv2d(in_channels=16, out_channels=8, kernel_size=7, padding=3),
|
||
# nn.ReLU(inplace=True),
|
||
# nn.BatchNorm2d(8),
|
||
)
|
||
#block2
|
||
self.layer5 = nn.Sequential(
|
||
nn.Conv2d(in_channels=3, out_channels=20, kernel_size=7, padding=3),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(20),
|
||
nn.MaxPool2d((2, 2))
|
||
)
|
||
self.layer6 = nn.Sequential(
|
||
nn.Conv2d(in_channels=20, out_channels=40, kernel_size=5, padding=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(40),
|
||
nn.MaxPool2d((2, 2))
|
||
)
|
||
self.layer7 = nn.Sequential(
|
||
nn.Conv2d(in_channels=40, out_channels=20, kernel_size=5, padding=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(20)
|
||
)
|
||
self.layer8 = nn.Sequential(
|
||
nn.Conv2d(in_channels=20, out_channels=10, kernel_size=5, padding=2),
|
||
# nn.ReLU(inplace=True),
|
||
# nn.BatchNorm2d(10)
|
||
)
|
||
#block3
|
||
self.layer9 = nn.Sequential(
|
||
nn.Conv2d(in_channels=3, out_channels=24, kernel_size=5, padding=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(24),
|
||
nn.MaxPool2d((2, 2))
|
||
)
|
||
self.layer10 = nn.Sequential(
|
||
nn.Conv2d(in_channels=24, out_channels=48, kernel_size=3, padding=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(48),
|
||
nn.MaxPool2d((2, 2))
|
||
)
|
||
self.layer11 = nn.Sequential(
|
||
nn.Conv2d(in_channels=48, out_channels=24, kernel_size=3, padding=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.BatchNorm2d(24)
|
||
)
|
||
self.layer12 = nn.Sequential(
|
||
nn.Conv2d(in_channels=24, out_channels=12, kernel_size=3, padding=1),
|
||
# nn.ReLU(inplace=True),
|
||
# nn.BatchNorm2d(12)
|
||
)
|
||
#end
|
||
self.layer13 = nn.Conv2d(in_channels=30, out_channels=1, kernel_size=1, padding=0)
|
||
|
||
|
||
def forward(self, x):
|
||
input = x.float()
|
||
x = self.layer1(input)
|
||
x = self.layer2(x)
|
||
x = self.layer3(x)
|
||
out1 = self.layer4(x)
|
||
|
||
x = self.layer5(input)
|
||
x = self.layer6(x)
|
||
x = self.layer7(x)
|
||
out2 = self.layer8(x)
|
||
|
||
x = self.layer9(input)
|
||
x = self.layer10(x)
|
||
x = self.layer11(x)
|
||
out3 = self.layer12(x)
|
||
|
||
out = torch.cat([out1,out2,out3], dim=1)
|
||
out = self.layer13(out)
|
||
|
||
return out
|
||
|
||
|
||
def _initialize_weights(self):
|
||
for m in self.modules():
|
||
if isinstance(m, nn.Conv2d):
|
||
nn.init.normal_(m.weight, std=0.01)
|
||
if m.bias is not None:
|
||
nn.init.constant_(m.bias, 0)
|
||
elif isinstance(m, nn.BatchNorm2d):
|
||
nn.init.constant_(m.weight, 1)
|
||
nn.init.constant_(m.bias, 0)
|
||
class PKSNet(nn.Module):
|
||
def __init__(self, load_vgg=True):
|
||
super(PKSNet, self).__init__()
|
||
#frontend_feat
|
||
self.frontend = nn.Sequential(
|
||
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True)
|
||
)
|
||
# left_header
|
||
self.left_header = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
|
||
nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
|
||
nn.AvgPool2d(kernel_size=2, stride=2),
|
||
nn.Flatten(),
|
||
nn.Linear(24576, 1000),
|
||
nn.Linear(1000, 5)
|
||
)
|
||
# right_header
|
||
self.right_header = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
|
||
nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
|
||
nn.AvgPool2d(kernel_size=2, stride=2),
|
||
nn.Flatten(),
|
||
nn.Linear(24576, 1000),
|
||
nn.Linear(1000, 13)
|
||
)
|
||
self.bacekend_feat = nn.Sequential(
|
||
#nn.Upsample(scale_factor=2, mode='nearest'),
|
||
#nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=(2 , 2), stride=2),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
#nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=(2 , 2), stride=2),
|
||
#nn.Upsample(scale_factor=2, mode='nearest'),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
#nn.ConvTranspose2d(in_channels=256, out_channels=256, kernel_size=(2 , 2), stride=2),
|
||
#nn.Upsample(scale_factor=2, mode='nearest'),
|
||
nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1)
|
||
)
|
||
|
||
if load_vgg:
|
||
mod = models.vgg16(pretrained=True)
|
||
self._initialize_weights()
|
||
keys1 = list(self.frontend.state_dict().keys())
|
||
keys2 = list(mod.state_dict().keys())
|
||
for i in range(len(keys1)):
|
||
self.frontend.state_dict()[keys1[i]].data[:] = mod.state_dict()[keys2[i]].data[:]
|
||
print('load vgg sucessfully!')
|
||
|
||
|
||
def forward(self, x):
|
||
x = x.float()
|
||
front_output = self.frontend(x)
|
||
# backend_feat
|
||
density = self.bacekend_feat(front_output)
|
||
# left_header
|
||
if model_head:
|
||
posture = self.left_header(front_output)
|
||
# right_header
|
||
spice = self.right_header(front_output)
|
||
temp = (density, posture, spice)
|
||
#return density, posture, spice
|
||
else:
|
||
temp = density
|
||
return temp
|
||
|
||
def _initialize_weights(self):
|
||
for m in self.modules():
|
||
if isinstance(m, nn.Conv2d):
|
||
nn.init.normal_(m.weight, std=0.01)
|
||
if m.bias is not None:
|
||
nn.init.constant_(m.bias, 0)
|
||
elif isinstance(m, nn.BatchNorm2d):
|
||
nn.init.constant_(m.weight, 1)
|
||
nn.init.constant_(m.bias, 0)
|
||
class PKSNet_attention(nn.Module):
|
||
def __init__(self, load_vgg=True):
|
||
super(PKSNet_attention, self).__init__()
|
||
#frontend_feat
|
||
self.frontend = nn.Sequential(
|
||
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True)
|
||
)
|
||
# left_header
|
||
self.left_header = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
|
||
)
|
||
self.left_header1 = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
|
||
nn.AvgPool2d(kernel_size=2, stride=2),
|
||
nn.Flatten(),
|
||
nn.Linear(24576, 1000),
|
||
nn.Linear(1000, 5)
|
||
)
|
||
# right_header
|
||
self.right_header = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
|
||
)
|
||
self.right_header1 = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
|
||
nn.AvgPool2d(kernel_size=2, stride=2),
|
||
nn.Flatten(),
|
||
nn.Linear(24576, 1000),
|
||
nn.Linear(1000, 13)
|
||
)
|
||
self.bacekend_feat = nn.Sequential(
|
||
#nn.Upsample(scale_factor=2, mode='nearest'),
|
||
#nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=(2 , 2), stride=2),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
#nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=(2 , 2), stride=2),
|
||
#nn.Upsample(scale_factor=2, mode='nearest'),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
#nn.ConvTranspose2d(in_channels=256, out_channels=256, kernel_size=(2 , 2), stride=2),
|
||
#nn.Upsample(scale_factor=2, mode='nearest'),
|
||
nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1)
|
||
)
|
||
|
||
if load_vgg:
|
||
mod = models.vgg16(pretrained=True)
|
||
self._initialize_weights()
|
||
keys1 = list(self.frontend.state_dict().keys())
|
||
keys2 = list(mod.state_dict().keys())
|
||
for i in range(len(keys1)):
|
||
self.frontend.state_dict()[keys1[i]].data[:] = mod.state_dict()[keys2[i]].data[:]
|
||
print('load vgg sucessfully!')
|
||
|
||
|
||
def forward(self, x):
|
||
x = x.float()
|
||
front_output = self.frontend(x)
|
||
# backend_feat
|
||
|
||
# left_header
|
||
if model_head:
|
||
posture_map = self.left_header(front_output)
|
||
|
||
posture = self.left_header1(posture_map)
|
||
# right_header
|
||
spice_map = self.right_header(front_output)
|
||
spice = self.right_header1(spice_map)
|
||
|
||
front_output_1 = front_output.clone()
|
||
#print(front_output)
|
||
with torch.set_grad_enabled(False):
|
||
for i in range(4):
|
||
for j in range(4):
|
||
#a = front_output[:,:,((i*24)):(((i+1)*24)),((j)*32):(((j+1)*32))]
|
||
front_output_1[:,:,((i*24)):(((i+1)*24)),((j)*32):(((j+1)*32))]=torch.mul(front_output[:,:,((i*24)):(((i+1)*24)),((j)*32):(((j+1)*32))],spice_map)
|
||
|
||
density = self.bacekend_feat(front_output_1)
|
||
temp = (density, posture, spice)
|
||
#print(front_output_1)
|
||
#return density, posture, spice
|
||
else:
|
||
density = self.bacekend_feat(front_output)
|
||
temp = density
|
||
return temp
|
||
|
||
def _initialize_weights(self):
|
||
for m in self.modules():
|
||
if isinstance(m, nn.Conv2d):
|
||
nn.init.normal_(m.weight, std=0.01)
|
||
if m.bias is not None:
|
||
nn.init.constant_(m.bias, 0)
|
||
elif isinstance(m, nn.BatchNorm2d):
|
||
nn.init.constant_(m.weight, 1)
|
||
nn.init.constant_(m.bias, 0)
|
||
class CSRNet(nn.Module):
|
||
def __init__(self, load_vgg=True):
|
||
super(MCNN, self).__init__()
|
||
#frontend_feat
|
||
self.frontend_feat = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]
|
||
self.backend_feat = [512, 512, 512, 256, 128, 64]
|
||
self.frontend = nn.Sequential(
|
||
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2, stride=2),
|
||
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
|
||
nn.ReLU(inplace=True)
|
||
)
|
||
# backend_feat
|
||
self.layer14 = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True))
|
||
self.layer15 = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True))
|
||
self.layer16 = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True))
|
||
self.layer17 = nn.Sequential(
|
||
nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True))
|
||
self.layer18 = nn.Sequential(
|
||
nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True))
|
||
self.layer19 = nn.Sequential(
|
||
nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=2, dilation=2),
|
||
nn.ReLU(inplace=True))
|
||
# output_layer
|
||
self.layer20 = nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1)
|
||
|
||
if load_vgg:
|
||
mod = models.vgg16(pretrained=True)
|
||
self._initialize_weights()
|
||
keys1 = list(self.frontend.state_dict().keys())
|
||
keys2 = list(mod.state_dict().keys())
|
||
for i in range(len(keys1)):
|
||
self.frontend.state_dict()[keys1[i]].data[:] = mod.state_dict()[keys2[i]].data[:]
|
||
print('load vgg sucessfully!')
|
||
|
||
|
||
def forward(self, x):
|
||
x = x.float()
|
||
x = self.frontend(x)
|
||
x = self.layer14(x)
|
||
x = self.layer15(x)
|
||
x = self.layer16(x)
|
||
x = self.layer17(x)
|
||
x = self.layer18(x)
|
||
x = self.layer19(x)
|
||
x = self.layer20(x)
|
||
return x
|
||
|
||
def _initialize_weights(self):
|
||
for m in self.modules():
|
||
if isinstance(m, nn.Conv2d):
|
||
nn.init.normal_(m.weight, std=0.01)
|
||
if m.bias is not None:
|
||
nn.init.constant_(m.bias, 0)
|
||
elif isinstance(m, nn.BatchNorm2d):
|
||
nn.init.constant_(m.weight, 1)
|
||
nn.init.constant_(m.bias, 0)
|
||
import torch.nn as nn
|
||
import torch
|
||
from torch.nn import functional as F
|
||
|
||
|
||
class ContextualModule(nn.Module):
|
||
def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
|
||
super(ContextualModule, self).__init__()
|
||
self.scales = []
|
||
self.scales = nn.ModuleList([self._make_scale(features, size) for size in sizes])
|
||
self.bottleneck = nn.Conv2d(features * 2, out_features, kernel_size=1)
|
||
self.relu = nn.ReLU()
|
||
self.weight_net = nn.Conv2d(features,features,kernel_size=1)
|
||
|
||
def __make_weight(self,feature,scale_feature):
|
||
weight_feature = feature - scale_feature
|
||
return F.sigmoid(self.weight_net(weight_feature))
|
||
|
||
def _make_scale(self, features, size):
|
||
prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
|
||
conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
|
||
return nn.Sequential(prior, conv)
|
||
|
||
def forward(self, feats):
|
||
h, w = feats.size(2), feats.size(3)
|
||
multi_scales = [F.upsample(input=stage(feats), size=(h, w), mode='bilinear') for stage in self.scales]
|
||
weights = [self.__make_weight(feats,scale_feature) for scale_feature in multi_scales]
|
||
overall_features = [(multi_scales[0]*weights[0]+multi_scales[1]*weights[1]+multi_scales[2]*weights[2]+multi_scales[3]*weights[3])/(weights[0]+weights[1]+weights[2]+weights[3])]+ [feats]
|
||
bottle = self.bottleneck(torch.cat(overall_features, 1))
|
||
return self.relu(bottle)
|
||
|
||
class CANNet(nn.Module):
|
||
def __init__(self, load_weights=False):
|
||
super(CANNet, self).__init__()
|
||
self.seen = 0
|
||
self.context = ContextualModule(512, 512)
|
||
self.frontend_feat = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]
|
||
self.backend_feat = [512, 512, 512,256,128,64]
|
||
self.frontend = make_layers(self.frontend_feat)
|
||
self.backend = make_layers(self.backend_feat,in_channels = 512,batch_norm=True, dilation = True)
|
||
self.output_layer = nn.Conv2d(64, 1, kernel_size=1)
|
||
|
||
|
||
def forward(self,x):
|
||
x = self.frontend(x)
|
||
x = self.context(x)
|
||
x = self.backend(x)
|
||
x = self.output_layer(x)
|
||
return x
|
||
|
||
def _initialize_weights(self):
|
||
for m in self.modules():
|
||
if isinstance(m, nn.Conv2d):
|
||
nn.init.normal_(m.weight, std=0.01)
|
||
if m.bias is not None:
|
||
nn.init.constant_(m.bias, 0)
|
||
elif isinstance(m, nn.BatchNorm2d):
|
||
nn.init.constant_(m.weight, 1)
|
||
nn.init.constant_(m.bias, 0)
|
||
|
||
def make_layers(cfg, in_channels = 3,batch_norm=False,dilation = False):
|
||
if dilation:
|
||
d_rate = 2
|
||
else:
|
||
d_rate = 1
|
||
layers = []
|
||
for v in cfg:
|
||
if v == 'M':
|
||
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
|
||
else:
|
||
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate,dilation = d_rate)
|
||
if batch_norm:
|
||
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
|
||
else:
|
||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||
in_channels = v
|
||
return nn.Sequential(*layers)
|
||
|
||
if __name__ == '__main__':
|
||
device = torch.device('cuda:2' if torch.cuda.is_available() else 'cpu')
|
||
print(device)
|
||
net = PKSNet_attention().to(device)
|
||
input = torch.randn(1,3,768,1024).to(device)
|
||
print(input)
|
||
density, posture, spice = net(input)
|
||
print(density.shape, posture.shape, spice.shape)
|