def get_model(model_name, pho_size=299, num_classes=110):
if model_name == "vgg16":
model = VGG(num_classes=num_classes, pho_size=299)
elif model_name == "resnet101":
model = resnet101(num_classes=num_classes)
elif model_name == "resnet152":
model = resnet152(num_classes=num_classes)
elif model_name == "densenet":
model = DenseNet(growth_rate=12,
block_config=[(100 - 4) // 6 for _ in range(3)],
num_classes=num_classes,
small_inputs=False,
efficient=True,
pho_size=pho_size)
elif model_name == "InceptionResNetV2":
model = InceptionResNetV2(num_classes=num_classes)
elif model_name == "InceptionV4":
model = InceptionV4(num_classes=num_classes)
elif model_name == "Inception3":
model = Inception3(num_classes=num_classes)
elif model_name == "denoise":
model = get_denoise()
elif model_name == "Mymodel":
model = Mymodel()
elif model_name == 'Comdefend':
model = ComDefend()
elif model_name == 'Rectifi':
model = Rectifi()
return model
def __init__(self,
embedding_size,
num_classes,
backbone='resnet18',
mode='t'):
super(background_resnet, self).__init__()
self.trainMode = mode
self.backbone = backbone
# copying modules from pretrained models
if backbone == 'resnet50':
self.pretrained = resnet.resnet50(pretrained=False)
elif backbone == 'resnet101':
self.pretrained = resnet.resnet101(pretrained=False)
elif backbone == 'resnet152':
self.pretrained = resnet.resnet152(pretrained=False)
elif backbone == 'resnet18':
self.pretrained = resnet.resnet18(pretrained=False)
elif backbone == 'resnet34':
self.pretrained = resnet.resnet34(pretrained=False)
else:
raise RuntimeError('unknown backbone: {}'.format(backbone))
self.fc0 = nn.Linear(128, embedding_size[0])
# task specific layers for task 1
self.fc1 = nn.Linear(128, embedding_size[1])
self.bn1 = nn.BatchNorm1d(embedding_size[1])
self.relu1 = nn.ReLU()
self.last1 = nn.Linear(embedding_size[1], num_classes)
# task speicific layers for task 2
self.fc2 = nn.Linear(128, embedding_size[2])
self.bn2 = nn.BatchNorm1d(embedding_size[2])
self.relu2 = nn.ReLU()
self.last2 = nn.Linear(embedding_size[2], num_classes)
def __init__(self, layers=18, bins=(1, 2, 3, 6), dropout=0.1, classes=2, zoom_factor=8, use_ppm=True,
criterion=nn.CrossEntropyLoss(ignore_index=255), BatchNorm=nn.BatchNorm2d, flow=False, sd=False,
pretrained=True):
super(PSPNet, self).__init__()
assert layers in [18, 50, 101, 152]
assert 512 % len(bins) == 0
assert classes > 1
assert zoom_factor in [1, 2, 4, 8]
self.zoom_factor = zoom_factor
self.use_ppm = use_ppm
self.flow = flow
self.sd = sd
self.criterion = criterion
models.BatchNorm = BatchNorm
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 18:
resnet = models.resnet18(deep_base=False, pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'conv1' in n:
m.stride = (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'conv1' in n:
m.stride = (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 512
if use_ppm:
self.ppm = PPM(fea_dim, int(fea_dim / len(bins)), bins, BatchNorm)
fea_dim *= 2
self.cls = nn.Sequential(
nn.Conv2d(fea_dim, 256, kernel_size=3, padding=1, bias=False),
BatchNorm(256),
nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout),
nn.Conv2d(256, classes, kernel_size=1)
)
if self.training:
self.aux = nn.Sequential(
nn.Conv2d(256, 256, kernel_size=3, padding=1, bias=False),
BatchNorm(256),
nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout),
nn.Conv2d(256, classes, kernel_size=1)
)
def __init__(self,
layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
classes=2,
zoom_factor=8,
use_ppm=True,
BatchNorm=nn.BatchNorm2d,
pretrained=False):
""" """
super(Backbone, self).__init__()
assert layers in [50, 101, 152]
assert 2048 % len(bins) == 0
assert classes > 1
assert zoom_factor in [1, 2, 4, 8]
self.zoom_factor = zoom_factor
self.use_ppm = use_ppm
models.BatchNorm = BatchNorm
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 2048
if use_ppm:
self.ppm = PPM(fea_dim, int(fea_dim / len(bins)), bins, BatchNorm)
fea_dim *= 2
self.cls = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
BatchNorm(512), nn.ReLU(inplace=True), nn.Dropout2d(p=dropout),
nn.Conv2d(512, classes, kernel_size=1))
if self.training:
self.aux = nn.Sequential(
nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
BatchNorm(256), nn.ReLU(inplace=True), nn.Dropout2d(p=dropout),
nn.Conv2d(256, classes, kernel_size=1))
def __init__(self):
super(Model, self).__init__()
resnet = resnet152(pretrained=True)
self.backbone = nn.Sequential(
resnet.conv1,
resnet.bn1,
resnet.relu,
resnet.maxpool,
resnet.layer1,
)
self.conv11 = resnet.layer_1
self.conv22 = resnet.layer_2
self.conv33 = resnet.layer_3
self.conv1_1 = nn.Sequential(nn.Conv2d(512, 512, 32),
nn.BatchNorm2d(512),
nn.LeakyReLU(True, inplace=0.2))
self.conv2_2 = nn.Sequential(nn.Conv2d(512, 512, 16),
nn.BatchNorm2d(512),
nn.LeakyReLU(True, inplace=0.2))
self.conv3_3 = nn.Sequential(nn.Conv2d(512, 512, 8),
nn.BatchNorm2d(512),
nn.LeakyReLU(True, inplace=0.2))
self.bottleneck_11 = nn.BatchNorm1d(512)
self.bottleneck_11.bias.requires_grad_(False)
self.classifier_11 = nn.Linear(in_features=512,
out_features=9999,
bias=False) # opt.class_num
self.color_11 = nn.Linear(in_features=512,
out_features=120,
bias=False)
self.bottleneck_22 = nn.BatchNorm1d(512)
self.bottleneck_22.bias.requires_grad_(False)
self.classifier_22 = nn.Linear(in_features=512,
out_features=9999,
bias=False)
self.color_22 = nn.Linear(in_features=512,
out_features=120,
bias=False)
self.bottleneck_33 = nn.BatchNorm1d(512)
self.bottleneck_33.bias.requires_grad_(False)
self.classifier_33 = nn.Linear(in_features=512,
out_features=9999,
bias=False)
self.color_33 = nn.Linear(in_features=512,
out_features=120,
bias=False)
def __init__(self, layers=50, dropout=0.1, classes=1, zoom_factor=8, criterion=nn.CrossEntropyLoss(ignore_index=255), BatchNorm=nn.BatchNorm2d, pretrained=True, args=None):
super(MGLNet, self).__init__()
assert layers in [50, 101, 152]
assert classes == 1
assert zoom_factor in [1, 2, 4, 8]
self.zoom_factor = zoom_factor
self.criterion = criterion
self.args = args
models.BatchNorm = BatchNorm
self.gamma = 1.0
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.conv2, resnet.bn2, resnet.relu, resnet.conv3, resnet.bn3, resnet.relu, resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
self.dim = 512
self.pred = nn.Sequential(
nn.Conv2d(2048, self.dim, kernel_size=3, padding=1, bias=False),
BatchNorm(self.dim),
nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout),
nn.Conv2d(self.dim, classes, kernel_size=1)
)
self.region_conv = self.pred[0:4] # 2048 -> 512
self.edge_cat = ConcatNet(BatchNorm) # concat low-level feature map to predict edge
# cascade mutual net
self.mutualnet0 = MutualNet(BatchNorm, dim=self.dim, num_clusters=args.num_clusters, dropout=dropout)
if args.stage == 1:
self.mutualnets = nn.ModuleList([self.mutualnet0])
elif args.stage == 2:
self.mutualnet1 = MutualNet(BatchNorm, dim=self.dim, num_clusters=args.num_clusters, dropout=dropout)
self.mutualnets = nn.ModuleList([self.mutualnet0, self.mutualnet1])
def __init__(self, embedding_size, num_classes, backbone='resnet50'):
super(background_resnet, self).__init__()
self.backbone = backbone
# copying modules from pretrained models
if backbone == 'resnet50':
self.pretrained = resnet.resnet50(pretrained=False)
elif backbone == 'resnet101':
self.pretrained = resnet.resnet101(pretrained=False)
elif backbone == 'resnet152':
self.pretrained = resnet.resnet152(pretrained=False)
elif backbone == 'resnet18':
self.pretrained = resnet.resnet18(pretrained=False)
elif backbone == 'resnet34':
self.pretrained = resnet.resnet34(pretrained=False)
else:
raise RuntimeError('unknown backbone: {}'.format(backbone))
self.fc0 = nn.Linear(512, embedding_size)
self.bn0 = nn.BatchNorm1d(embedding_size)
self.relu = nn.ReLU()
self.last = nn.Linear(embedding_size, num_classes)
def get_model(model_type='resnet50', num_classes=1000):
# TODO: Add more backbones
if model_type == 'resnet34':
model = resnet.resnet34(pretrained=True)
elif model_type == 'resnet50':
model = resnet.resnet50(pretrained=True)
elif model_type == 'resnet101':
model = resnet.resnet101(pretrained=True)
elif model_type == 'resnet152':
model = resnet.resnet152(pretrained=True)
elif model_type == 'resnext50_32x4d':
model = resnet.resnext50_32x4d(pretrained=True)
elif model_type == 'resnext101_32x8d':
model = resnet.resnext101_32x8d(pretrained=True)
elif model_type == 'res2net_v1b_50':
model = res2net50_v1b_26w_4s(pretrained=True)
elif model_type == 'res2net_v1b_101':
model = res2net101_v1b_26w_4s(pretrained=True)
elif model_type == 'res2net50_26w_4s':
model = res2net50_26w_4s(pretrained=True)
elif model_type == 'res2net101_26w_4s':
model = res2net101_26w_4s(pretrained=True)
elif model_type == 'res2next50':
model = res2next50(pretrained=True)
elif model_type == 'senet154':
model = senet.senet154(num_classes=num_classes, pretrained='imagenet')
elif model_type == 'resnest50':
model = resnest50(pretrained=True)
elif model_type == 'resnest101':
model = resnest101(pretrained=True)
elif model_type == 'resnest200':
model = resnest200(pretrained=True)
elif model_type == 'resnest269':
model = resnest269(pretrained=True)
else:
model = resnet.resnet50(pretrained=True)
return model
def __init__(self,
layers=50,
dropout=0.1,
classes=2,
zoom_factor=8,
criterion=nn.CrossEntropyLoss(ignore_index=255),
BatchNorm=nn.BatchNorm2d,
pretrained=True):
super(DeepLabV3, self).__init__()
print(layers)
assert layers in [50, 101, 152]
assert classes > 1
assert zoom_factor in [1, 2, 4, 8]
self.zoom_factor = zoom_factor
self.criterion = criterion
models.BatchNorm = BatchNorm
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
grids = [1, 2, 4]
countConvolutions = 0
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (
2 * grids[countConvolutions], 2 *
grids[countConvolutions]), (2 * grids[countConvolutions],
2 *
grids[countConvolutions]), (1,
1)
countConvolutions = countConvolutions + 1
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 2048
pyramids = [6, 12, 18]
self.aspp = _ASPPModule(2048, 256, pyramids)
self.cls = nn.Sequential(
_ConvBatchNormReLU(256 * (len(pyramids) + 2), 256, 1, 1, 0, 1),
nn.Conv2d(256, classes, kernel_size=1))
if self.training:
self.aux = nn.Sequential(
nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
BatchNorm(256), nn.ReLU(inplace=True), nn.Dropout2d(p=dropout),
nn.Conv2d(256, classes, kernel_size=1))
def __init__(self,
layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
classes=2,
zoom_factor=8,
use_ppm=True,
criterion=nn.CrossEntropyLoss(ignore_index=255),
BatchNorm=nn.BatchNorm2d,
pretrained=True):
super(PSPNet, self).__init__()
assert layers in [18, 50, 101, 152]
assert 2048 % len(bins) == 0
assert classes > 1
assert zoom_factor in [1, 2, 4, 8]
self.zoom_factor = zoom_factor
self.use_ppm = use_ppm
self.criterion = criterion
self.criterion_reg = nn.MSELoss(reduce=False)
models.BatchNorm = BatchNorm
if layers == 18:
resnet = models_origin.resnet18(pretrained=True)
elif layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
'''self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 512
if use_ppm:
self.ppm = PPM(fea_dim, int(fea_dim/len(bins)), bins, BatchNorm)
fea_dim *= 2'''
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 2048
if use_ppm:
self.ppm = PPM(fea_dim, int(fea_dim / len(bins)), bins, BatchNorm)
fea_dim *= 2
'''self.cls = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
BatchNorm(512),
nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout),
nn.Conv2d(512, classes, kernel_size=1)
)'''
self.cls3 = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
BatchNorm(512), nn.ReLU(inplace=True),
nn.Conv2d(512, 96 * 2, kernel_size=1), BatchNorm(96 * 2),
nn.ReLU(inplace=True))
self.cls2 = nn.Sequential(
nn.Conv2d(96 * 2 + 96, 96, kernel_size=3, padding=1, bias=False),
BatchNorm(96), nn.ReLU(inplace=True), nn.Dropout2d(p=dropout),
nn.Conv2d(96, classes, kernel_size=1))
self.reg = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
BatchNorm(512), nn.ReLU(inplace=True),
nn.Conv2d(512, 96, kernel_size=1))
self.reg2 = nn.Sequential(
nn.Conv2d(96, 96, kernel_size=3, padding=1, bias=False),
BatchNorm(96), nn.ReLU(inplace=True),
nn.Conv2d(96, 96, kernel_size=1), nn.ReLU(inplace=True))
self.reg3 = nn.Sequential(
nn.Conv2d(96, 96, kernel_size=3, padding=1, bias=False),
BatchNorm(96), nn.ReLU(inplace=True),
nn.Conv2d(96, 96, kernel_size=1), nn.ReLU(inplace=True))
self.bn = BatchNorm(96)
self.bn2 = BatchNorm(96)
self.mean = torch.tensor(np.load('../data/meanvar/mean.npy'),
requires_grad=False)
self.var = torch.tensor(np.sqrt(np.load('../data/meanvar/var.npy')),
requires_grad=False)
self.mean_2d = torch.tensor(np.load('../data/meanvar/mean_2d.npy'),
requires_grad=False) #.cuda().float()
self.var_2d = torch.tensor(np.sqrt(
np.load('../data/meanvar/var_2d.npy')),
requires_grad=False) #.cuda().float()
self.var = self.var[0, :]
self.mean = torch.unsqueeze(self.mean, 0)
self.mean = torch.unsqueeze(self.mean, 2)
self.mean = torch.unsqueeze(self.mean, 2)
self.mean = self.mean.repeat(1, 1, 60, 60)
self.var = torch.unsqueeze(self.var, 0)
self.var = torch.unsqueeze(self.var, 2)
self.var = torch.unsqueeze(self.var, 2)
self.var = self.var.repeat(1, 1, 60, 60)
elif args.layer == '50':
net = resnet50(pretrained=False,
progress=True,
activate=activate,
hidden=hidden,
num_classes=10)
elif args.layer == '101':
net = resnet101(pretrained=False,
progress=True,
activate=activate,
hidden=hidden,
num_classes=10)
elif args.layer == '152':
net = resnet152(pretrained=False,
progress=True,
activate=activate,
hidden=hidden,
num_classes=10)
else:
raise ValueError('layer should be 18 / 34 / 50 / 101 / 152')
net.to(device)
net.initialize_weights(zero_init_residual=True)
# ============================ step 3/6 损失函数 ============================
criterion = nn.CrossEntropyLoss() # 选择损失函数
# ============================ step 4/6 优化器 ============================
# 选择优化器
if args.optim == 'sgd':
optimizer = optim.SGD(net.parameters(),
def get_edge_scores(junctions,
regions,
rgb_folder,
_id,
epoch=1,
model='resnet152'):
# load model
resnet = resnet152(pretrained=False).cuda()
edge_classifier = EdgeClassifier(resnet)
edge_classifier = edge_classifier.cuda()
edge_classifier = edge_classifier.eval()
# open RGB image
split = 'det'
out_size = 256
rgb_path = os.path.join(rgb_folder, _id + '.jpg')
rgb = Image.open(rgb_path).resize((out_size, out_size))
rgb = np.array(rgb) / 255.0
model_path = '/home/nelson/Workspace/building_reconstruction/working_model/binary_edge_classifier_with_regions/saved_models/edge_classifier_{}_{}_iter_{}.pth'.format(
model, split, epoch)
edge_classifier.load_state_dict(torch.load(model_path))
# check backup -- save time
temp_dir = './temp/{}/'.format(model)
if not os.path.isdir(temp_dir):
os.makedirs(temp_dir)
backup_path = '{}/{}_{}_{}.npy'.format(temp_dir, _id, epoch, split)
if os.path.isfile(backup_path):
return np.load(open(backup_path, 'rb'), encoding='bytes').item()
# combine regions
all_reg = np.zeros((out_size, out_size))
for k, reg in enumerate(regions):
reg = Image.fromarray(reg * 255.0).resize((out_size, out_size))
reg = np.array(reg) / 255.0
inds = np.array(np.where(reg == 1))
all_reg[inds[0, :], inds[1, :]] = k
# draw edge
lw_from_cls = {}
for k, c1 in enumerate(junctions):
for l, c2 in enumerate(junctions):
if k > l:
edge = Image.new('L', (out_size, out_size))
dr = ImageDraw.Draw(edge)
x1, y1 = c1
x2, y2 = c2
div = 256.0 / out_size
dr.line((int(x1 / div), int(y1 / div), int(
x2 / div), int(y2 / div)),
fill="white",
width=int(4 / div))
edge = np.array(edge) / 255.0
imgs = np.concatenate([
rgb.transpose(2, 0, 1), edge[np.newaxis, :, :],
all_reg[np.newaxis, :, :]
], 0)
imgs = torch.from_numpy(imgs).cuda().float()
prob = edge_classifier(imgs.unsqueeze(0))
prob = prob.detach().cpu().numpy()
lw_from_cls[(k, l)] = prob[0]
# save backup
np.save(open(backup_path, 'wb'), lw_from_cls)
return lw_from_cls
writer = None
# batch size
train_batchSize = [args.label_batch_size, args.unlabel_batch_size]
# backbone architecture
if args.arch == 'resnet18':
backbone = resnet.resnet18(feature_len=args.feat_len)
elif args.arch == 'resnet34':
backbone = resnet.resnet34(feature_len=args.feat_len)
elif args.arch == 'resnet50':
backbone = resnet.resnet50(feature_len=args.feat_len)
elif args.arch == 'resnet101':
backbone = resnet.resnet101(feature_len=args.feat_len)
elif args.arch == 'resnet152':
backbone = resnet.resnet152(feature_len=args.feat_len)
elif args.arch == 'usr':
backbone = model_usr
else:
raise NameError(
'Arch %s is not support. Please enter from [resnet18, resnet34, resnet50, resnet101, resnet152, usr]'
% args.arch)
# head
model_head = arc.ArcMarginProduct_virface(in_features=args.feat_len,
out_features=args.num_ids,
s=32,
m=0.5,
device='cuda')
# generator, model
if args.method == 'pretrain':
def __init__(self, layers=50, classes=2, zoom_factor=8, \
criterion=nn.CrossEntropyLoss(ignore_index=255), BatchNorm=nn.BatchNorm2d, \
pretrained=True, sync_bn=True, shot=1, ppm_scales=[60, 30, 15, 8], vgg=False):
super(PFENet, self).__init__()
assert layers in [50, 101, 152]
print(ppm_scales)
assert classes > 1
from torch.nn import BatchNorm2d as BatchNorm
self.zoom_factor = zoom_factor
self.criterion = criterion
self.shot = shot
self.ppm_scales = ppm_scales
self.vgg = vgg
models.BatchNorm = BatchNorm
if self.vgg:
print('INFO: Using VGG_16 bn')
vgg_models.BatchNorm = BatchNorm
vgg16 = vgg_models.vgg16_bn(pretrained=pretrained)
print(vgg16)
self.layer0, self.layer1, self.layer2, \
self.layer3, self.layer4 = get_vgg16_layer(vgg16)
else:
print('INFO: Using ResNet {}'.format(layers))
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu1,
resnet.conv2, resnet.bn2, resnet.relu2,
resnet.conv3, resnet.bn3, resnet.relu3,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
reduce_dim = 256
if self.vgg:
fea_dim = 512 + 256
else:
fea_dim = 1024 + 512
self.cls = nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True), nn.Dropout2d(p=0.1),
nn.Conv2d(reduce_dim, classes, kernel_size=1))
self.down_query = nn.Sequential(
nn.Conv2d(fea_dim,
reduce_dim,
kernel_size=1,
padding=0,
bias=False), nn.ReLU(inplace=True), nn.Dropout2d(p=0.5))
self.down_supp = nn.Sequential(
nn.Conv2d(fea_dim,
reduce_dim,
kernel_size=1,
padding=0,
bias=False), nn.ReLU(inplace=True), nn.Dropout2d(p=0.5))
self.pyramid_bins = ppm_scales
self.avgpool_list = []
for bin in self.pyramid_bins:
if bin > 1:
self.avgpool_list.append(nn.AdaptiveAvgPool2d(bin))
factor = 1
mask_add_num = 1
self.init_merge = []
self.beta_conv = []
self.inner_cls = []
for bin in self.pyramid_bins:
self.init_merge.append(
nn.Sequential(
nn.Conv2d(reduce_dim * 2 + mask_add_num,
reduce_dim,
kernel_size=1,
padding=0,
bias=False),
nn.ReLU(inplace=True),
))
self.beta_conv.append(
nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True)))
self.inner_cls.append(
nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Dropout2d(p=0.1),
nn.Conv2d(reduce_dim, classes, kernel_size=1)))
self.init_merge = nn.ModuleList(self.init_merge)
self.beta_conv = nn.ModuleList(self.beta_conv)
self.inner_cls = nn.ModuleList(self.inner_cls)
self.res1 = nn.Sequential(
nn.Conv2d(reduce_dim * len(self.pyramid_bins),
reduce_dim,
kernel_size=1,
padding=0,
bias=False),
nn.ReLU(inplace=True),
)
self.res2 = nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False),
nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False),
nn.ReLU(inplace=True),
)
self.GAP = nn.AdaptiveAvgPool2d(1)
self.alpha_conv = []
for idx in range(len(self.pyramid_bins) - 1):
self.alpha_conv.append(
nn.Sequential(
nn.Conv2d(512,
256,
kernel_size=1,
stride=1,
padding=0,
bias=False), nn.ReLU()))
self.alpha_conv = nn.ModuleList(self.alpha_conv)
def __init__(self,
layers=50,
bins=(1, 3, 6, 8),
dropout=0.2,
classes=2,
zoom_factor=8,
use_aspp=True,
output_stride=8,
criterion=nn.CrossEntropyLoss(ignore_index=255),
hidden_dim=512,
BatchNorm=nn.BatchNorm2d,
pretrained=True,
ImLength=110):
super(TransformNet, self).__init__()
assert layers in [50, 101, 152]
assert 2048 % len(bins) == 0
assert classes > 1
assert zoom_factor in [1, 2, 4, 8]
self.zoom_factor = zoom_factor
self.use_aspp = use_aspp
self.criterion = criterion
self.os = output_stride
self.bins = bins
models.BatchNorm = BatchNorm
if layers == 50:
resnet = models.resnet50(output_stride,
BatchNorm,
pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(output_stride,
BatchNorm,
pretrained=pretrained)
else:
resnet = models.resnet152(output_stride,
BatchNorm,
pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
self.feat_proj = nn.Conv2d(resnet.num_channels[-1],
hidden_dim,
kernel_size=1)
self.ppm = PSPModule(sizes=(1, 3, 6, 8), dimension=2)
self.pos_enc = PositionEmbeddingSine(hidden_dim // 2, normalize=True)
self.transformer = Transformer(d_model=hidden_dim,
nhead=8,
num_encoder_layers=0,
num_decoder_layers=6,
dim_feedforward=2048,
dropout=0.1,
activation="relu",
normalize_before=False,
return_intermediate_dec=False)
self.query_embed = nn.Embedding(ImLength, hidden_dim)
self.cls = nn.Sequential(
nn.Conv2d(hidden_dim * (1 + len(bins)),
512,
kernel_size=3,
padding=1,
bias=False), BatchNorm(512), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout), nn.Conv2d(512, classes, kernel_size=1))
if self.training:
self.aux = nn.Sequential(
nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
BatchNorm(256), nn.ReLU(inplace=True), nn.Dropout2d(p=dropout),
nn.Conv2d(256, classes, kernel_size=1))
self.tr_dec_aux1 = nn.Sequential(
nn.Conv2d(hidden_dim * (1 + len(bins)),
512,
kernel_size=3,
padding=1,
bias=False), BatchNorm(512), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout), nn.Conv2d(512, classes,
kernel_size=1))
def __init__(self, args):
super(Model, self).__init__()
layers = args.layers
classes = args.classes
sync_bn = args.sync_bn
pretrained = True
assert layers in [50, 101, 152]
assert classes > 1
from torch.nn import BatchNorm2d as BatchNorm
self.zoom_factor = args.zoom_factor
self.criterion = nn.CrossEntropyLoss(ignore_index=255)
self.shot = args.shot
self.train_iter = args.train_iter
self.eval_iter = args.eval_iter
self.pyramid = args.pyramid
models.BatchNorm = BatchNorm
print('INFO: Using ResNet {}'.format(layers))
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu1,
resnet.conv2, resnet.bn2, resnet.relu2,
resnet.conv3, resnet.bn3, resnet.relu3,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
reduce_dim = 256
fea_dim = 1024 + 512
self.cls = nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True), nn.Dropout2d(p=0.1),
nn.Conv2d(reduce_dim, classes, kernel_size=1))
self.down_conv = nn.Sequential(
nn.Conv2d(fea_dim,
reduce_dim,
kernel_size=1,
padding=0,
bias=False), nn.Dropout2d(p=0.5))
# Using Feature Enrichment Module from PFENet as context module
if self.pyramid:
self.pyramid_bins = args.ppm_scales
self.avgpool_list = []
for bin in self.pyramid_bins:
if bin > 1:
self.avgpool_list.append(nn.AdaptiveAvgPool2d(bin))
self.corr_conv = []
self.beta_conv = []
self.inner_cls = []
for bin in self.pyramid_bins:
self.corr_conv.append(
nn.Sequential(
nn.Conv2d(reduce_dim * 2 + 1,
reduce_dim,
kernel_size=1,
padding=0,
bias=False),
nn.ReLU(inplace=True),
))
self.beta_conv.append(
nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True)))
self.inner_cls.append(
nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Dropout2d(p=0.1),
nn.Conv2d(reduce_dim, classes, kernel_size=1)))
self.corr_conv = nn.ModuleList(self.corr_conv)
self.beta_conv = nn.ModuleList(self.beta_conv)
self.inner_cls = nn.ModuleList(self.inner_cls)
self.alpha_conv = []
for idx in range(len(self.pyramid_bins) - 1):
self.alpha_conv.append(
nn.Sequential(
nn.Conv2d(2 * reduce_dim,
reduce_dim,
kernel_size=1,
stride=1,
padding=0,
bias=False), nn.ReLU(inplace=True)))
self.alpha_conv = nn.ModuleList(self.alpha_conv)
self.res1 = nn.Sequential(
nn.Conv2d(reduce_dim * len(self.pyramid_bins),
reduce_dim,
kernel_size=1,
padding=0,
bias=False),
nn.ReLU(inplace=True),
)
self.res2 = nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False),
nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False),
nn.ReLU(inplace=True),
)
# Using ASPP as context module
else:
self.ASPP = ASPP(out_channels=reduce_dim)
self.corr_conv = nn.Sequential(
nn.Conv2d(reduce_dim * 2 + 1,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Dropout2d(p=0.5))
self.skip1 = nn.Sequential(
nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
stride=1,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
stride=1,
padding=1,
bias=False))
self.skip2 = nn.Sequential(
nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
stride=1,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
stride=1,
padding=1,
bias=False))
self.skip3 = nn.Sequential(
nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
stride=1,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
stride=1,
padding=1,
bias=False))
self.decoder = build_decoder(256)
self.cls_aux = nn.Sequential(
nn.Conv2d(reduce_dim,
reduce_dim,
kernel_size=3,
padding=1,
bias=False), nn.ReLU(inplace=True),
nn.Dropout2d(p=0.1),
nn.Conv2d(reduce_dim, classes, kernel_size=1))
parser.add_argument("--record_path", type=str)
args = parser.parse_args()
######## Choose whether to use Cuda
device = torch.device("cuda" if (
args.use_cuda and torch.cuda.is_available()) else "cpu")
torch.cuda.empty_cache()
####### Pick according model
if args.model_name == 'res50':
model = resnet.resnet50().to(device)
elif args.model_name == 'res101':
model = resnet.resnet101().to(device)
elif args.model_name == 'res152':
model = resnet.resnet152().to(device)
elif args.model_name == 'res34':
model = resnet.resnet34().to(device)
elif args.model_name == 'res18':
model = resnet.resnet18().to(device)
elif args.model_name == 'alexnet':
model = alexnet.alexnet().to(device)
else:
print('Wrong Model Name')
########### Whether to parallel the model
if args.use_cuda:
if args.parallel:
model = nn.DataParallel(model)
else:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_num
def main():
parser = argparse.ArgumentParser("PyTorch Face Recognizer")
parser.add_argument('--cmd',
default='test',
type=str,
choices=['train', 'test', 'extract'],
help='train, test or extract')
parser.add_argument('--arch_type',
type=str,
default='resnet50_ft',
help='model type',
choices=[
'resnet50_ft', 'senet50_ft', 'resnet50_scratch',
'senet50_scratch', 'resnet152'
])
parser.add_argument('--dataset_dir',
type=str,
default='/media/hyo/文档/Dataset/vggface2/train',
help='dataset directory')
parser.add_argument('--log_file',
type=str,
default='./log_file',
help='log file')
parser.add_argument(
'--train_img_list_file',
type=str,
default='/media/hyo/文档/Dataset/vggface2/train_list.txt',
help='text file containing image files used for training')
parser.add_argument(
'--test_img_list_file',
type=str,
default='/media/hyo/文档/Dataset/vggface2/train_list.txt',
help=
'text file containing image files used for validation, test or feature extraction'
)
parser.add_argument(
'--meta_file',
type=str,
default='/media/hyo/文档/Dataset/vggface2/identity_meta.csv',
help='meta file')
parser.add_argument('--checkpoint_dir',
type=str,
default='./checkpoint',
help='checkpoints directory')
parser.add_argument('--feature_dir',
type=str,
default='./feature',
help='directory where extracted features are saved')
parser.add_argument(
'-c',
'--config',
type=int,
default=1,
choices=configurations.keys(),
help='the number of settings and hyperparameters used in training')
parser.add_argument('--batch_size',
type=int,
default=32,
help='batch size')
parser.add_argument('--resume',
type=str,
default='',
help='checkpoint file')
parser.add_argument('--weight_file',
type=str,
default='./weight/weight.pkl',
help='weight file')
parser.add_argument('--gpu', type=int, default=0)
parser.add_argument('-j',
'--workers',
default=4,
type=int,
metavar='N',
help='number of data loading workers (default: 4)')
parser.add_argument(
'--horizontal_flip',
default=True,
action='store_true',
help='horizontally flip images specified in test_img_list_file')
args = parser.parse_args()
print(args)
if args.cmd == "extract":
utils.create_dir(args.feature_dir)
if args.cmd == 'train':
utils.create_dir(args.checkpoint_dir)
cfg = configurations[args.config]
log_file = args.log_file
resume = args.resume
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
cuda = torch.cuda.is_available()
if cuda:
print("torch.backends.cudnn.version: {}".format(
torch.backends.cudnn.version()))
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
# 0. id label map
meta_file = args.meta_file
id_label_dict = utils.get_id_label_map(meta_file)
weight_file = args.weight_file
# 1. data loader
root = args.dataset_dir
train_img_list_file = args.train_img_list_file
test_img_list_file = args.test_img_list_file
kwargs = {'num_workers': args.workers, 'pin_memory': True} if cuda else {}
if args.cmd == 'train':
dt = datasets.VGG_Faces2(root,
train_img_list_file,
id_label_dict,
split='train')
train_loader = torch.utils.data.DataLoader(dt,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
dv = datasets.VGG_Faces2(root,
test_img_list_file,
id_label_dict,
split='valid',
horizontal_flip=args.horizontal_flip)
val_loader = torch.utils.data.DataLoader(dv,
batch_size=args.batch_size,
shuffle=False,
**kwargs)
# 2. model
include_top = True if args.cmd != 'extract' else False
if 'resnet' in args.arch_type:
model = ResNet.resnet152(num_classes=N_IDENTITY,
include_top=include_top)
# else:
# model = SENet.senet50(num_classes=N_IDENTITY, include_top=include_top)
print(model)
start_epoch = 0
start_iteration = 0
if resume:
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
start_iteration = checkpoint['iteration']
assert checkpoint['arch'] == args.arch_type
print("Resume from epoch: {}, iteration: {}".format(
start_epoch, start_iteration))
# else:
# utils.load_state_dict(model, args.weight_file)
if args.cmd == 'train':
model.fc.reset_parameters()
if cuda:
model = model.cuda()
criterion = nn.CrossEntropyLoss()
if cuda:
criterion = criterion.cuda()
# 3. optimizer
if args.cmd == 'train':
optim = torch.optim.SGD([
{
'params': get_parameters(model, bias=False)
},
{
'params': get_parameters(model, bias=True),
'lr': cfg['lr'] * 2,
'weight_decay': 0
},
],
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
if resume:
optim.load_state_dict(checkpoint['optim_state_dict'])
# lr_policy: step
last_epoch = start_iteration if resume else -1
lr_scheduler = torch.optim.lr_scheduler.StepLR(optim,
cfg['step_size'],
gamma=cfg['gamma'],
last_epoch=last_epoch)
if args.cmd == 'train':
trainer = Trainer(
cmd=args.cmd,
cuda=cuda,
model=model,
criterion=criterion,
optimizer=optim,
lr_scheduler=lr_scheduler,
train_loader=train_loader,
val_loader=val_loader,
log_file=log_file,
max_iter=cfg['max_iteration'],
checkpoint_dir=args.checkpoint_dir,
print_freq=100,
model_dict=weight_file,
)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
elif args.cmd == 'test':
validator = Validator(
cmd=args.cmd,
cuda=cuda,
model=model,
criterion=criterion,
val_loader=val_loader,
log_file=log_file,
print_freq=100,
)
validator.validate()
elif args.cmd == 'extract':
extractor = Extractor(
cuda=cuda,
model=model,
val_loader=val_loader,
log_file=log_file,
feature_dir=args.feature_dir,
flatten_feature=True,
print_freq=100,
)
extractor.extract()
def fine_tune(load_dir,
save_dir,
filename,
dataloader_train,
dataloader_val,
ContinuousTrain=False):
'''
load pretrained model
resnet-18-kinetics.pth: --model resnet --model_depth 18 --resnet_shortcut A
resnet-34-kinetics.pth: --model resnet --model_depth 34 --resnet_shortcut A
resnet-34-kinetics-cpu.pth: CPU ver. of resnet-34-kinetics.pth
resnet-50-kinetics.pth: --model resnet --model_depth 50 --resnet_shortcut B
resnet-101-kinetics.pth: --model resnet --model_depth 101 --resnet_shortcut B
resnet-152-kinetics.pth: --model resnet --model_depth 152 --resnet_shortcut B
resnet-200-kinetics.pth: --model resnet --model_depth 200 --resnet_shortcut B
preresnet-200-kinetics.pth: --model preresnet --model_depth 200 --resnet_shortcut B
wideresnet-50-kinetics.pth: --model wideresnet --model_depth 50 --resnet_shortcut B --wide_resnet_k 2
resnext-101-kinetics.pth: --model resnext --model_depth 101 --resnet_shortcut B --resnext_cardinality 32
densenet-121-kinetics.pth: --model densenet --model_depth 121
densenet-201-kinetics.pth: --model densenet --model_depth 201
'''
num_epochs = 100
step = 0
if not ContinuousTrain:
model = resnet.resnet152(sample_size=112,
sample_duration=args.n_frames_per_clip,
shortcut_type='B',
num_classes=83)
checkpoint = utils.load_checkpoint(load_dir, filename)
model = nn.DataParallel(model, device_ids=[1])
state_dict = deepcopy(model.state_dict())
feature_state_dict = {
key: value
for key, value in checkpoint['state_dict'].items()
if key not in ['module.fc.weight', 'module.fc.bias']
}
state_dict.update(feature_state_dict)
model.load_state_dict(state_dict)
# set fine tune parameters: Conv5_x and fc layer from original paper
for param in model.module.parameters():
param.requires_grad = False
for named_child in model.module.named_children():
if named_child[0] == 'fc' or named_child[0] == 'layer4':
# if named_child[0] == 'fc':
for param in named_child[1].parameters():
param.requires_grad = True
# pdb.set_trace()
else:
print('Recover model from {}/resnet18.pth......'.format(save_dir))
checkpoint = utils.load_checkpoint(save_dir, 'resnet18.pth')
model = checkpoint['model']
state_dict = checkpoint['state_dict']
model.load_state_dict(state_dict)
step = checkpoint['step']
model.to(device)
summary(model, (3, args.n_frames_per_clip, 112, 112))
# pdb.set_trace()
model.train()
# determine optimizer
criterion = nn.CrossEntropyLoss()
fc_lr_layers = list(map(id, model.module.fc.parameters()))
pretrained_lr_layers = [
p for p in model.parameters()
if id(p) not in fc_lr_layers and p.requires_grad == True
]
# pretrained_lr_layers = filter(lambda p:
# id(p) not in fc_lr_layers, model.parameters())
optimizer = torch.optim.SGD([{
"params": model.module.fc.parameters()
}, {
"params": pretrained_lr_layers,
"lr": 1e-4,
'weight_decay': 1e-3
}],
lr=1e-3,
momentum=0.9,
weight_decay=1e-2)
train_logger = utils.Logger(
os.path.join('output', 'R3D-fine-tune-all.log'), [
'step', 'train_loss', 'train_acc', 'val_loss', 'val_acc',
'lr_feature', 'lr_fc'
])
scheduler = lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)
# scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=2)
train_loss = utils.AverageMeter()
train_acc = utils.AverageMeter()
val_loss = utils.AverageMeter()
val_acc = utils.AverageMeter()
for epoch in trange(num_epochs): # loop over the dataset multiple times
train_loss.reset()
train_acc.reset()
for data in dataloader_train:
inputs, masks, labels = data
inputs, labels = inputs.to(device,
non_blocking=True).float(), labels.to(
device, non_blocking=True).long()
optimizer.zero_grad()
outputs = model(inputs)
loss_ = criterion(outputs, labels)
loss_.backward()
optimizer.step()
train_loss.update(loss_.item())
train_acc.update(utils.calculate_accuracy(outputs, labels))
if step % 50 == 0:
val_loss.reset()
val_acc.reset()
model.eval()
for data_val in dataloader_val:
inputs_val, masks_val, labels_val = data_val
inputs_val, labels_val = inputs_val.to(device, non_blocking=True).float(), \
labels_val.to(device, non_blocking=True).long()
outputs_val = model(inputs_val)
val_loss_ = criterion(outputs_val, labels_val)
val_loss.update(val_loss_.item())
val_acc.update(
utils.calculate_accuracy(outputs_val, labels_val))
model.train()
print(
'epoch{}/{} train_acc:{:.3f} train_loss:{:.3f} val_acc:{:.3f} val_loss:{:.3f}'
.format(epoch + 1, num_epochs, train_acc.val,
train_loss.val, val_acc.avg, val_loss.avg))
train_logger.log({
'step': step,
'train_loss': train_loss.val,
'train_acc': train_acc.val,
'val_loss': val_loss.avg,
'val_acc': val_acc.avg,
# 'lr_feature': optimizer.param_groups[1]['lr'],
'lr_feature': 0,
'lr_fc': optimizer.param_groups[0]['lr']
})
if step % 200 == 0:
utils.save_checkpoint(model, optimizer, step, save_dir,
'resnet18.pth')
step += 1
scheduler.step()
def __init__(self,
tag='teacher',
layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
classes=19,
zoom_factor=8,
use_ppm=True,
criterion=nn.CrossEntropyLoss(ignore_index=255),
pretrained=True):
super(PSPNet, self).__init__()
assert layers in [18, 50, 101, 152]
assert 2048 % len(bins) == 0
assert classes > 1
assert zoom_factor in [1, 2, 4, 8]
self.zoom_factor = zoom_factor
self.use_ppm = use_ppm
self.criterion = criterion
self.tag = tag
if layers == 18:
resnet = models.resnet18(pretrained=pretrained)
elif layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
if layers in [50, 101]:
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
fea_dim = 2048
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
elif layers in [18]:
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
fea_dim = 512
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
if 'downsample.0' in n:
m.stride = (1, 1)
if 'conv1' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
if 'downsample.0' in n:
m.stride = (1, 1)
if 'conv1' in n:
m.stride = (1, 1)
if use_ppm:
self.ppm = PPM(fea_dim, int(fea_dim / len(bins)), bins)
fea_dim *= 2
self.cls = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout), nn.Conv2d(512, classes, kernel_size=1))
if self.training and self.tag == 'student':
self.aux = nn.Sequential(
nn.Conv2d(256, 128, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(128), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout), nn.Conv2d(128, classes,
kernel_size=1))
def __init__(self,
layers=50,
dropout=0.1,
classes=2,
zoom_factor=8,
use_psa=True,
psa_type=2,
compact=False,
shrink_factor=2,
mask_h=59,
mask_w=59,
normalization_factor=1.0,
psa_softmax=True,
criterion=nn.CrossEntropyLoss(ignore_index=255),
BatchNorm=nn.BatchNorm2d,
pretrained=True):
super(PSANet, self).__init__()
assert layers in [50, 101, 152]
assert classes > 1
assert zoom_factor in [1, 2, 4, 8]
assert psa_type in [0, 1, 2]
self.zoom_factor = zoom_factor
self.use_psa = use_psa
self.criterion = criterion
models.BatchNorm = BatchNorm
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 2048
if use_psa:
self.psa = PSA(fea_dim, 512, psa_type, compact, shrink_factor,
mask_h, mask_w, normalization_factor, psa_softmax,
BatchNorm)
fea_dim *= 2
self.cls = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
BatchNorm(512), nn.ReLU(inplace=True), nn.Dropout2d(p=dropout),
nn.Conv2d(512, classes, kernel_size=1))
if self.training:
self.aux = nn.Sequential(
nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
BatchNorm(256), nn.ReLU(inplace=True), nn.Dropout2d(p=dropout),
nn.Conv2d(256, classes, kernel_size=1))
def __init__(self, layers=50, classes=2, criterion=nn.CrossEntropyLoss(ignore_index=255),
pretrained=True, shot=1, ppm_scales=[60, 30, 15, 8], vgg=False, FPN=True):
super(FSSNet, self).__init__()
assert layers in [50, 101, 152]
print(ppm_scales)
assert classes > 1
from torch.nn import BatchNorm2d as BatchNorm
# 参数
self.criterion = criterion
self.shot = shot
self.vgg = vgg
self.trip = nn.Parameter(torch.zeros(1))
models.BatchNorm = BatchNorm
# Backbone Related
if self.vgg:
print('INFO: Using VGG_16 bn')
vgg_models.BatchNorm = BatchNorm
vgg16 = vgg_models.vgg16_bn(pretrained=pretrained)
print(vgg16)
self.layer0, self.layer1, self.layer2, \
self.layer3, self.layer4 = get_vgg16_layer(vgg16)
else:
print('INFO: Using ResNet {}'.format(layers))
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu1, resnet.conv2, resnet.bn2, resnet.relu2,
resnet.conv3, resnet.bn3, resnet.relu3, resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
# feature dimension
reduce_dim = 256
class_num =2
# if self.vgg:
# fea_dim = 512 + 256
# else:
# fea_dim = 1024 + 512
# query 下采样
self.down = nn.Sequential(
nn.Conv2d(512, reduce_dim, kernel_size=1, padding=0, bias=False),
nn.ReLU(inplace=True),
nn.Dropout2d(p=0.5)
)
self.down_s = nn.Sequential(
nn.Conv2d(1024, reduce_dim, kernel_size=1, padding=0, bias=False),
nn.ReLU(inplace=True),
nn.Dropout2d(p=0.5)
)
#
# # support 下采样
self.Tripletencoder = Tripletencoder101(FPN).cuda()
# 256 -> 256
self.res1 = nn.Sequential(
nn.Conv2d(reduce_dim, reduce_dim, kernel_size=1, padding=0, bias=False),
nn.ReLU(inplace=True),
)
# 256 -> 256
self.res2 = nn.Sequential(
nn.Conv2d(reduce_dim, reduce_dim, kernel_size=3, padding=1, bias=False),
nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim, reduce_dim, kernel_size=3, padding=1, bias=False),
nn.ReLU(inplace=True),
)
# 256 -> 256
self.res3 = nn.Sequential(
nn.Conv2d(reduce_dim, reduce_dim, kernel_size=3, padding=1, bias=False),
nn.ReLU(inplace=True),
nn.Conv2d(reduce_dim, reduce_dim, kernel_size=3, padding=1, bias=False),
nn.ReLU(inplace=True),
)
# ASPP
self.ASPP = ASPP().cuda()
self.GCN = g_GCN(reduce_dim, int(reduce_dim / 2)).cuda()
# 512 -> 256
self.cls = nn.Sequential(
nn.Conv2d(reduce_dim, reduce_dim, kernel_size=3, padding=1, bias=False),
nn.ReLU(inplace=True),
nn.Dropout2d(p=0.1),
nn.Conv2d(reduce_dim, classes, kernel_size=1)
)
self.int1 = nn.Sequential(
nn.Conv2d(reduce_dim+1, reduce_dim, kernel_size=1, padding=0, bias=False),
nn.ReLU(inplace=True))
self.int2 = nn.Sequential(
nn.Conv2d(reduce_dim+1, reduce_dim, kernel_size=1, padding=0, bias=False),
nn.ReLU(inplace=True))
self.int3 = nn.Sequential(
nn.Conv2d(reduce_dim+1, reduce_dim, kernel_size=1, padding=0, bias=False),
nn.ReLU(inplace=True))
