forked from wffjwbbf/ComDesignProject
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import torch
|
|
from torch import nn
|
|
import torch.nn.functional as F
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
# Define model
|
|
# input:128x256x3
|
|
# output: 64
|
|
class ReIDNet(nn.Module):
|
|
def __init__(self):
|
|
super(ReIDNet, 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 |