forked from wffjwbbf/ComDesignProject
229 lines
7.3 KiB
Python
229 lines
7.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class BasicBlock(nn.Module):
|
|
def __init__(self, c_in, c_out,is_downsample=False):
|
|
super(BasicBlock,self).__init__()
|
|
self.is_downsample = is_downsample
|
|
if is_downsample:
|
|
self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=2, padding=1, bias=False)
|
|
else:
|
|
self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=1, padding=1, bias=False)
|
|
self.bn1 = nn.BatchNorm2d(c_out)
|
|
self.relu = nn.ReLU(True)
|
|
self.conv2 = nn.Conv2d(c_out,c_out,3,stride=1,padding=1, bias=False)
|
|
self.bn2 = nn.BatchNorm2d(c_out)
|
|
if is_downsample:
|
|
self.downsample = nn.Sequential(
|
|
nn.Conv2d(c_in, c_out, 1, stride=2, bias=False),
|
|
nn.BatchNorm2d(c_out)
|
|
)
|
|
elif c_in != c_out:
|
|
self.downsample = nn.Sequential(
|
|
nn.Conv2d(c_in, c_out, 1, stride=1, bias=False),
|
|
nn.BatchNorm2d(c_out)
|
|
)
|
|
self.is_downsample = True
|
|
|
|
def forward(self,x):
|
|
y = self.conv1(x)
|
|
y = self.bn1(y)
|
|
y = self.relu(y)
|
|
y = self.conv2(y)
|
|
y = self.bn2(y)
|
|
if self.is_downsample:
|
|
x = self.downsample(x)
|
|
return F.relu(x.add(y),True)
|
|
|
|
def make_layers(c_in,c_out,repeat_times, is_downsample=False):
|
|
blocks = []
|
|
for i in range(repeat_times):
|
|
if i ==0:
|
|
blocks += [BasicBlock(c_in,c_out, is_downsample=is_downsample),]
|
|
else:
|
|
blocks += [BasicBlock(c_out,c_out),]
|
|
return nn.Sequential(*blocks)
|
|
|
|
class Net(nn.Module):
|
|
def __init__(self, num_classes=751 ,reid=False):
|
|
super(Net,self).__init__()
|
|
# 3 128 64
|
|
self.conv = nn.Sequential(
|
|
nn.Conv2d(3,64,3,stride=1,padding=1),
|
|
nn.BatchNorm2d(64),
|
|
nn.ReLU(inplace=True),
|
|
# nn.Conv2d(32,32,3,stride=1,padding=1),
|
|
# nn.BatchNorm2d(32),
|
|
# nn.ReLU(inplace=True),
|
|
nn.MaxPool2d(3,2,padding=1),
|
|
)
|
|
# 32 64 32
|
|
self.layer1 = make_layers(64,64,2,False)
|
|
# 32 64 32
|
|
self.layer2 = make_layers(64,128,2,True)
|
|
# 64 32 16
|
|
self.layer3 = make_layers(128,256,2,True)
|
|
# 128 16 8
|
|
self.layer4 = make_layers(256,512,2,True)
|
|
# 256 8 4
|
|
self.avgpool = nn.AvgPool2d((8,4),1)
|
|
# 256 1 1
|
|
self.reid = reid
|
|
self.classifier = nn.Sequential(
|
|
nn.Linear(512, 256),
|
|
nn.BatchNorm1d(256),
|
|
nn.ReLU(inplace=True),
|
|
nn.Dropout(),
|
|
nn.Linear(256, num_classes),
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
x = self.layer1(x)
|
|
x = self.layer2(x)
|
|
x = self.layer3(x)
|
|
x = self.layer4(x)
|
|
x = self.avgpool(x)
|
|
x = x.view(x.size(0),-1)
|
|
# B x 128
|
|
if self.reid:
|
|
x = x.div(x.norm(p=2,dim=1,keepdim=True))
|
|
return x
|
|
# classifier
|
|
x = self.classifier(x)
|
|
return x
|
|
|
|
# Define model
|
|
# input:128x256x3
|
|
# output: 64
|
|
class MyNet(nn.Module):
|
|
def __init__(self):
|
|
super(MyNet, self).__init__()
|
|
self.conv1 = nn.Conv2d(3,8,5,1,2)
|
|
self.maxpool1 = nn.MaxPool2d(2) # 64x128x8
|
|
self.bn1 = nn.BatchNorm2d(8)
|
|
self.conv2 = nn.Conv2d(8,16,5,1,2)
|
|
self.maxpool2 = nn.MaxPool2d(2) # 32x64x16
|
|
self.bn2 = nn.BatchNorm2d(16)
|
|
self.conv3 = nn.Conv2d(16,32,5,1,2)
|
|
self.conv4 = nn.Conv2d(32,64,5,1,2)
|
|
self.maxpool3 = nn.MaxPool2d(2) # 16x32x64
|
|
self.bn3 = nn.BatchNorm2d(64)
|
|
self.conv5 = nn.Conv2d(64,32,5,1,2) # 16x32x32
|
|
self.maxpool4 = nn.MaxPool2d(2) # 8x16x32
|
|
self.bn4 = nn.BatchNorm2d(32)
|
|
self.conv6 = nn.Conv2d(32,64,8,8) #1x2x64
|
|
self.flat = nn.Flatten()
|
|
self.linear = nn.Linear(2*64,64) # 64D-feature ID
|
|
|
|
self.sigmoid = nn.Sigmoid()
|
|
self.lrelu = nn.LeakyReLU()
|
|
|
|
def forward(self, x:torch.Tensor):
|
|
x = self.conv1(x)
|
|
x = self.lrelu(x)
|
|
x = self.maxpool1(x)
|
|
x = self.bn1(x)
|
|
x = self.conv2(x)
|
|
x = self.lrelu(x)
|
|
x = self.maxpool2(x)
|
|
x = self.bn2(x)
|
|
x = self.conv3(x)
|
|
x = self.lrelu(x)
|
|
x = self.conv4(x)
|
|
x = self.lrelu(x)
|
|
x = self.maxpool3(x)
|
|
x = self.bn3(x)
|
|
x = self.conv5(x)
|
|
x = self.lrelu(x)
|
|
x = self.maxpool4(x)
|
|
x = self.bn4(x)
|
|
x = self.conv6(x)
|
|
x = self.sigmoid(x)
|
|
x = self.flat(x)
|
|
x = self.linear(x)
|
|
|
|
return x
|
|
|
|
if __name__ == '__main__':
|
|
net = Net()
|
|
x = torch.randn(4,3,128,64)
|
|
y = net(x)
|
|
|
|
#########ResNet18
|
|
class RestNetBasicBlock(nn.Module):
|
|
def __init__(self, in_channels, out_channels, stride):
|
|
super(RestNetBasicBlock, self).__init__()
|
|
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
|
|
self.bn1 = nn.BatchNorm2d(out_channels)
|
|
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
|
|
self.bn2 = nn.BatchNorm2d(out_channels)
|
|
|
|
def forward(self, x):
|
|
output = self.conv1(x)
|
|
output = F.relu(self.bn1(output))
|
|
output = self.conv2(output)
|
|
output = self.bn2(output)
|
|
return F.leaky_relu(x + output)
|
|
|
|
|
|
class RestNetDownBlock(nn.Module):
|
|
def __init__(self, in_channels, out_channels, stride):
|
|
super(RestNetDownBlock, self).__init__()
|
|
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride[0], padding=1)
|
|
self.bn1 = nn.BatchNorm2d(out_channels)
|
|
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride[1], padding=1)
|
|
self.bn2 = nn.BatchNorm2d(out_channels)
|
|
self.extra = nn.Sequential(
|
|
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride[0], padding=0),
|
|
nn.BatchNorm2d(out_channels)
|
|
)
|
|
|
|
def forward(self, x):
|
|
extra_x = self.extra(x)
|
|
output = self.conv1(x)
|
|
out = F.relu(self.bn1(output))
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out)
|
|
return F.leaky_relu(extra_x + out)
|
|
|
|
|
|
class RestNet18(nn.Module):
|
|
def __init__(self):
|
|
super(RestNet18, self).__init__()
|
|
self.conv1 = nn.Conv2d(3, 32, kernel_size=(3, 3), padding=1)
|
|
self.bn1 = nn.BatchNorm2d(32)
|
|
self.maxpool = nn.MaxPool2d(2)
|
|
|
|
self.layer1 = nn.Sequential(RestNetBasicBlock(32, 32, 1),
|
|
RestNetBasicBlock(32, 32, 1))
|
|
|
|
self.layer2 = nn.Sequential(RestNetDownBlock(32, 64, [2, 1]),
|
|
RestNetBasicBlock(64, 64, 1))
|
|
|
|
self.layer3 = nn.Sequential(RestNetDownBlock(64, 128, [2, 1]),
|
|
RestNetBasicBlock(128, 128, 1))
|
|
|
|
self.layer4 = nn.Sequential(RestNetDownBlock(128, 256, [2, 1]),
|
|
RestNetBasicBlock(256, 256, 1))
|
|
|
|
self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
|
|
|
|
self.fc = nn.Linear(32768, 23)
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
x = self.layer1(x)
|
|
x = self.layer2(x)
|
|
x = self.layer3(x)
|
|
x = self.layer4(x)
|
|
# x = self.avgpool(x)
|
|
x = x.reshape(x.shape[0], -1)
|
|
# print(x.size())
|
|
x = self.fc(x)
|
|
return x
|
|
|
|
|