def __init__(self):
super(YangNet, self).__init__()
self.config = Config()
self.conv1 = nn.Sequential( # input shape(1, 96, 96)
nn.Conv2d(1, 96, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(96),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(96, 48, 48)
self.conv2 = nn.Sequential(
nn.Conv2d(96, 128, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(128, 24, 24)
self.conv3 = nn.Sequential(
nn.Conv2d(128, 160, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(160),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(160, 12, 12)
self.conv4 = nn.Sequential(
nn.Conv2d(160, 256, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(256, 6, 6)
self.conv5 = nn.Sequential(
nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(384),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(384, 3, 3)
self.conv6 = nn.Sequential(
nn.Conv2d(384, 1024, kernel_size=1, stride=1, padding=1,
bias=True),
nn.BatchNorm2d(1024),
nn.ReLU(),
) # output shape(1024, 3, 3)
self.conv7 = nn.Sequential(
nn.Conv2d(1024,
1024,
kernel_size=1,
stride=1,
padding=1,
bias=True), nn.BatchNorm2d(1024), nn.ReLU(),
nn.AdaptiveAvgPool2d(1)) # output shape(1024, 1, 1)
self.fc = nn.Sequential(
nn.Linear(1024, self.config.random_size, bias=True), )
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.xavier_normal_(m.weight.data)
def __init__(self):
super(AlexNet_SE, self).__init__()
self.config = Config()
self.conv1 = nn.Sequential( # input shape(1, 114, 114)
nn.Conv2d(1, 96, kernel_size=11, stride=4, bias=True),
nn.BatchNorm2d(96),
nn.ReLU(),
nn.MaxPool2d(kernel_size=3, stride=2),
) # output shape(96, 13, 13)
self.se1 = SELayer(channel=96, reduction=16)
self.conv2 = nn.Sequential(
nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2, bias=True),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(256, 6, 6)
self.se2 = SELayer(channel=256, reduction=16)
self.conv3 = nn.Sequential(
nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(384),
nn.ReLU(),
) # output shape(384, 6, 6)
self.se3 = SELayer(channel=384, reduction=16)
self.conv4 = nn.Sequential(
nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(384),
nn.ReLU(),
) # output shape(384, 6, 6)
self.se4 = SELayer(channel=384, reduction=16)
self.conv5 = nn.Sequential(
nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(256, 3, 3)
self.fc = nn.Sequential(
nn.Dropout(p=0.5),
nn.Linear(256*3*3, 4096, bias=True),
nn.BatchNorm1d(4096),
nn.ReLU(),
nn.Dropout(p=0.5),
nn.Linear(4096, 4096, bias=True),
nn.BatchNorm1d(4096),
nn.ReLU(),
nn.Linear(4096, self.config.random_size, bias=True),
)
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.xavier_normal_(m.weight.data)
sys.path.append('../../')
from models.YangNet import YangNet
from training.AlexNet.trainer import AlexNetTrainer
from training.AlexNet.config import Config
from tools.logger import Logger
from checkpoint import CheckPoint
import os
import tools.imagedb as imagedb
import time
import multiprocessing
if __name__ == '__main__':
multiprocessing.freeze_support()
config = Config()
if not os.path.exists(config.save_path):
os.makedirs(config.save_path)
os.environ['CUDA VISIBLE_DEVICES'] = config.GPU
use_cuda = config.use_cuda and torch.cuda.is_available()
torch.manual_seed(config.manualSeed)
torch.cuda.manual_seed(config.manualSeed)
device = torch.device("cuda" if use_cuda else "cpu")
torch.backends.cudnn.benchmark = True
imagedb.preprocess_gnt()
train_annotation_path = os.path.join(config.dataPath,
'train_png_annotation.txt')
valid_annotation_path = os.path.join(config.dataPath,
'valid_png_annotation.txt')
def __init__(self):
super(AlexNet_ST, self).__init__()
self.config = Config()
self.conv1 = nn.Sequential( # input shape(1, 114, 114)
nn.Conv2d(1, 96, kernel_size=11, stride=4, bias=True),
nn.BatchNorm2d(96),
nn.ReLU(),
nn.MaxPool2d(kernel_size=3, stride=2),
) # output shape(96, 13, 13)
self.conv2 = nn.Sequential(
nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2, bias=True),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(256, 6, 6)
self.conv3 = nn.Sequential(
nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(384),
nn.ReLU(),
) # output shape(384, 6, 6)
self.conv4 = nn.Sequential(
nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(384),
nn.ReLU(),
) # output shape(384, 6, 6)
self.conv5 = nn.Sequential(
nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1, bias=True),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
) # output shape(256, 3, 3)
self.fc = nn.Sequential(
nn.Dropout(p=0.5),
nn.Linear(256 * 3 * 3, 4096, bias=True),
nn.BatchNorm1d(4096),
nn.ReLU(),
nn.Dropout(p=0.5),
nn.Linear(4096, 4096, bias=True),
nn.BatchNorm1d(4096),
nn.ReLU(),
nn.Linear(4096, self.config.random_size, bias=True),
)
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.xavier_normal_(m.weight.data)
# Spatial transformer localization-network
self.localization = nn.Sequential(
nn.Conv2d(
1,
16,
kernel_size=11,
stride=4,
), # input shape(1, 114, 114)
nn.BatchNorm2d(16),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True),
nn.Conv2d(16, 32, kernel_size=5), # input shape(16, 13, 13)
nn.BatchNorm2d(32),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True),
nn.Conv2d(32, 64, kernel_size=3,
padding=1), # input shape(32, 4, 4)
nn.AvgPool2d((4, 4)),
nn.ReLU(True),
)
# Regressor for the 3 * 2 affine matrix
self.fc_loc = nn.Sequential(nn.Linear(64, 32), nn.ReLU(True),
nn.Linear(32, 3 * 2))
# Initialize the weights/bias with identity transformation
self.fc_loc[2].weight.data.zero_()
self.fc_loc[2].bias.data.copy_(
torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))