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))