class BFE(nn.Module):
def __init__(self, num_classes=CATEGORIES, width_ratio=0.5, height_ratio=0.5, model_path=None):
super(BFE, self).__init__()
resnet = resnet50(pretrained=True)
self.backbone = nn.Sequential(
resnet.conv1,
resnet.bn1,
resnet.relu,
resnet.maxpool,
resnet.layer1, # res_conv2
resnet.layer2, # res_conv3
resnet.layer3, # res_conv4
)
self.res_part = nn.Sequential(
Bottleneck(1024, 512, stride=1, downsample=nn.Sequential(
nn.Conv2d(1024, 2048, kernel_size=1, stride=1, bias=False),
nn.BatchNorm2d(2048),
)),
Bottleneck(2048, 512),
Bottleneck(2048, 512),
)
self.res_part.load_state_dict(resnet.layer4.state_dict())
reduction = nn.Sequential(
nn.Conv2d(2048, 512, 1),
nn.BatchNorm2d(512),
nn.ReLU()
)
# global branch
self.global_avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.global_softmax = nn.Linear(512, num_classes)
self.global_softmax.apply(weights_init_kaiming)
self.global_reduction = copy.deepcopy(reduction)
self.global_reduction.apply(weights_init_kaiming)
# part branch
self.res_part2 = Bottleneck(2048, 512)
self.part_maxpool = nn.AdaptiveMaxPool2d((1,1))
self.batch_crop = BatchDrop(height_ratio, width_ratio)
self.reduction = nn.Sequential(
nn.Linear(2048, 1024, 1),
nn.BatchNorm1d(1024),
nn.ReLU()
)
self.reduction.apply(weights_init_kaiming)
self.softmax = nn.Linear(1024, num_classes)
self.softmax.apply(weights_init_kaiming)
state = load_model(model_path)
if state:
new_state = self.state_dict()
new_state.update({k: v for k, v in state.items() if k in new_state})
self.load_state_dict(new_state)
def forward(self, x):
"""
:param x: input image tensor of (N, C, H, W)
:return: (prediction, triplet_losses, softmax_losses)
"""
x = self.backbone(x)
x = self.res_part(x)
predict = []
triplet_features = []
softmax_features = []
#global branch
glob = self.global_avgpool(x)
global_triplet_feature = self.global_reduction(glob).squeeze()
global_softmax_class = self.global_softmax(global_triplet_feature)
softmax_features.append(global_softmax_class)
triplet_features.append(global_triplet_feature)
predict.append(global_triplet_feature)
#part branch
x = self.res_part2(x)
x = self.batch_crop(x)
triplet_feature = self.part_maxpool(x).squeeze()
feature = self.reduction(triplet_feature)
softmax_feature = self.softmax(feature)
triplet_features.append(feature)
softmax_features.append(softmax_feature)
predict.append(feature)
if self.training:
return triplet_features, softmax_features
else:
return torch.cat(predict, 1)
def get_optim_policy(self):
params = [
{'params': self.backbone.parameters()},
{'params': self.res_part.parameters()},
{'params': self.global_reduction.parameters()},
{'params': self.global_softmax.parameters()},
{'params': self.res_part2.parameters()},
{'params': self.reduction.parameters()},
{'params': self.softmax.parameters()},
]
return params
class BFE(nn.Module):
def __init__(self, num_classes, width_ratio=0.5, height_ratio=0.5):
super(BFE, self).__init__()
resnet = resnet50(pretrained=True)
self.backbone = nn.Sequential(
resnet.conv1,
resnet.bn1,
resnet.relu,
resnet.maxpool,
resnet.layer1, # res_conv2
resnet.layer2, # res_conv3
resnet.layer3, # res_conv4
)
self.res_part = nn.Sequential(
Bottleneck(1024,
512,
stride=1,
downsample=nn.Sequential(
nn.Conv2d(1024,
2048,
kernel_size=1,
stride=1,
bias=False),
nn.BatchNorm2d(2048),
)),
Bottleneck(2048, 512),
Bottleneck(2048, 512),
)
self.res_part.load_state_dict(resnet.layer4.state_dict())
reduction = nn.Sequential(nn.Conv2d(2048, 512, 1), nn.BatchNorm2d(512),
nn.ReLU())
# stage 1
self.global_avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.global_softmax = nn.Linear(512, num_classes)
self.global_softmax.apply(weights_init_kaiming)
self.global_reduction = copy.deepcopy(reduction)
self.global_reduction.apply(weights_init_kaiming)
# stage 2
self.res_part2 = Bottleneck(2048, 512)
self.part_maxpool = nn.AdaptiveMaxPool2d((1, 1))
# stage 2
self.cutmix_batch_drop1 = CutMixBatchDrop(height_ratio, width_ratio)
self.reduction1 = nn.Sequential(nn.Linear(2048, 512, 1),
nn.BatchNorm1d(512), nn.ReLU())
self.reduction1.apply(weights_init_kaiming)
self.softmax1 = nn.Linear(512, num_classes)
self.softmax1.apply(weights_init_kaiming)
# stage 3
self.cutmix_batch_drop2 = CutMixBatchDrop(height_ratio, width_ratio)
self.reduction2 = nn.Sequential(nn.Linear(2048, 512, 1),
nn.BatchNorm1d(512), nn.ReLU())
self.reduction2.apply(weights_init_kaiming)
self.softmax2 = nn.Linear(512, num_classes)
self.softmax2.apply(weights_init_kaiming)
def forward(self, x, y=None):
"""
:param x: input image tensor of (N, C, H, W)
:return: (prediction, triplet_losses, softmax_losses)
"""
x = self.backbone(x)
x = self.res_part(x)
predict = []
triplet_features = []
softmax_features = []
# stage 1
glob = self.global_avgpool(x)
global_triplet_feature = self.global_reduction(glob).squeeze()
global_softmax_class = self.global_softmax(global_triplet_feature)
softmax_features.append(global_softmax_class)
triplet_features.append(global_triplet_feature)
predict.append(global_triplet_feature)
# stage 2
x_part = self.res_part2(x)
x_part1 = x_part
x_part2 = x_part
# stage 2
if self.training:
gradcam, _ = gradCAM(x, global_softmax_class, y)
x_part1, idx_cutmix1, lamda = self.cutmix_batch_drop1(
x_part1, gradcam, y, [3, 2], 3)
cutmix1_triplet_feature = self.part_maxpool(
x_part1).squeeze() # N*2048
cutmix1_feature = self.reduction1(cutmix1_triplet_feature) # N*1024
cutmix1_softmax_feature = self.softmax1(
cutmix1_feature) # N*num_class/751
triplet_features.append(cutmix1_feature)
softmax_features.append(cutmix1_softmax_feature)
predict.append(cutmix1_feature)
# stage 3
if self.training:
gradcam, _ = gradCAM(x_part1, cutmix1_softmax_feature, y)
x_part2, idx_cutmix2, lamda = self.cutmix_batch_drop1(
x_part2, gradcam, y, [3, 2], 3)
cutmix2_triplet_feature = self.part_maxpool(
x_part2).squeeze() # N*2048
cutmix2_feature = self.reduction2(cutmix2_triplet_feature) # N*1024
cutmix2_softmax_feature = self.softmax2(
cutmix2_feature) # N*num_class/751
triplet_features.append(cutmix2_feature)
softmax_features.append(cutmix2_softmax_feature)
predict.append(cutmix2_feature)
if self.training:
return triplet_features, softmax_features
else:
return torch.cat(predict, 1)
def get_optim_policy(self):
params = [
{
'params': self.backbone.parameters()
},
{
'params': self.res_part.parameters()
},
{
'params': self.global_reduction.parameters()
},
{
'params': self.global_softmax.parameters()
},
{
'params': self.res_part2.parameters()
},
{
'params': self.reduction1.parameters()
},
{
'params': self.softmax1.parameters()
},
{
'params': self.reduction2.parameters()
},
{
'params': self.softmax2.parameters()
},
]
return params
class BFE(nn.Module):
def __init__(self, num_classes, width_ratio=0.5, height_ratio=0.5):
super(BFE, self).__init__()
# resnet = resnet50()
# resnet.load_state_dict(torch.load('./resnet50-19c8e357.pth'))
# print("load ResNet50 parameters")
print("load ResNeSt50 and ResNet50 parameters")
model = resnet50()
resnet = resnest50()
model.load_state_dict(torch.load('../resnet50-19c8e357.pth'))
resnet.load_state_dict(torch.load('../resnest50-528c19ca.pth'))
self.backbone = nn.Sequential(
resnet.conv1,
resnet.bn1,
resnet.relu,
resnet.maxpool,
resnet.layer1,
CAM_Module(256),
resnet.layer2,
CAM_Module(512),
)
self.res_part_head = nn.Sequential(resnet.layer3, CAM_Module(1024))
self.res_part = nn.Sequential(
Bottleneck(1024,
512,
stride=1,
downsample=nn.Sequential(
nn.Conv2d(1024,
2048,
kernel_size=1,
stride=1,
bias=False),
nn.BatchNorm2d(2048),
)),
Bottleneck(2048, 512),
Bottleneck(2048, 512),
)
# self.res_part.load_state_dict(resnet.layer4.state_dict())
self.res_part.load_state_dict(model.layer4.state_dict())
self.c4 = CAM_Module(2048)
self.part_pool = nn.AdaptiveMaxPool2d((1, 1))
# self.part_pool = nn.AdaptiveAvgPool2d((1, 1))
self.res_part2 = Bottleneck(2048, 512)
self.batch_crop = BatchDrop(height_ratio, width_ratio)
self.reduction = nn.Sequential(nn.Linear(2048, 1024, 1),
nn.BatchNorm1d(1024), nn.ReLU())
self.reduction.apply(weights_init_kaiming)
self.softmax = nn.Linear(1024, num_classes)
self.softmax.apply(weights_init_kaiming)
# Auxiliary branch
ab_vector_size = 256
reduction = nn.Sequential(nn.Conv2d(2048, ab_vector_size, 1),
nn.BatchNorm2d(ab_vector_size),
nn.ReLU(inplace=True))
self.auxiliary_avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.auxiliary_softmax = nn.Linear(ab_vector_size, num_classes)
self.auxiliary_softmax.apply(weights_init_kaiming)
self.auxiliary_reduction = copy.deepcopy(reduction)
self.auxiliary_reduction.apply(weights_init_kaiming)
def forward(self, x):
x = self.backbone(x)
x = self.res_part_head(x)
x = self.res_part(x)
x = self.c4(x)
predict = []
triplet_features = []
softmax_features = []
# auxiliary branch
auxiliary = self.auxiliary_avgpool(x)
auxiliary_triplet_feature = self.auxiliary_reduction(
auxiliary).squeeze()
auxiliary_softmax_class = self.auxiliary_softmax(
auxiliary_triplet_feature)
softmax_features.append(auxiliary_softmax_class)
triplet_features.append(auxiliary_triplet_feature)
predict.append(auxiliary_triplet_feature)
# main branch
x = self.res_part2(x)
x = self.batch_crop(x)
triplet_feature = self.part_pool(x).squeeze()
feature = self.reduction(triplet_feature)
softmax_feature = self.softmax(feature)
triplet_features.append(feature)
softmax_features.append(softmax_feature)
predict.append(feature)
if self.training:
return triplet_features, softmax_features
else:
return torch.cat(predict, 1)
def get_optim_policy(self):
params = [
{
'params': self.backbone.parameters()
},
{
'params': self.res_part_head.parameters()
},
{
'params': self.res_part.parameters()
},
{
'params': self.c4.parameters()
},
{
'params': self.res_part2.parameters()
},
{
'params': self.reduction.parameters()
},
{
'params': self.softmax.parameters()
},
{
'params': self.auxiliary_reduction.parameters()
},
{
'params': self.auxiliary_softmax.parameters()
},
]
return params
class BFE(nn.Module):
def __init__(self, num_classes, width_ratio=0.5, height_ratio=0.5):
super(BFE, self).__init__()
resnet = resnet50(pretrained=True)
self.backbone = nn.Sequential(
resnet.conv1,
resnet.bn1,
resnet.relu,
resnet.maxpool,
resnet.layer1, # res_conv2
resnet.layer2, # res_conv3
resnet.layer3, # res_conv4
)
self.res_part = nn.Sequential(
Bottleneck(1024,
512,
stride=1,
downsample=nn.Sequential(
nn.Conv2d(1024,
2048,
kernel_size=1,
stride=1,
bias=False),
nn.BatchNorm2d(2048),
)),
Bottleneck(2048, 512),
Bottleneck(2048, 512),
)
self.res_part.load_state_dict(resnet.layer4.state_dict())
reduction = nn.Sequential(nn.Conv2d(2048, 512, 1), nn.BatchNorm2d(512),
nn.ReLU())
# global branch
self.global_avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.global_softmax = nn.Linear(512, num_classes)
self.global_softmax.apply(weights_init_kaiming)
self.global_reduction = copy.deepcopy(reduction)
self.global_reduction.apply(weights_init_kaiming)
self.global_conv_list = nn.ModuleList()
for _ in range(6):
self.global_conv_list.append(
nn.Sequential(nn.Conv2d(2048, 256, 1), nn.BatchNorm2d(256),
nn.ReLU()))
self.global_softmax_list = nn.ModuleList()
for _ in range(6):
self.global_softmax_list.append(
nn.Sequential(nn.Linear(256, num_classes)))
for i in range(6):
self.global_conv_list[i].apply(weights_init_kaiming)
self.global_softmax_list[i].apply(weights_init_kaiming)
# part branch
self.res_part2 = Bottleneck(2048, 512)
self.part_maxpool = nn.AdaptiveMaxPool2d((1, 1))
self.batch_crop = BatchDrop(height_ratio, width_ratio)
self.reduction = nn.Sequential(nn.Linear(2048, 1024, 1),
nn.BatchNorm1d(1024), nn.ReLU())
self.reduction.apply(weights_init_kaiming)
self.softmax = nn.Linear(1024, num_classes)
self.softmax.apply(weights_init_kaiming)
def forward(self, x):
"""
:param x: input image tensor of (N, C, H, W)
:return: (prediction, triplet_losses, softmax_losses)
"""
x = self.backbone(x)
x = self.res_part(x)
predict = []
triplet_features = []
softmax_features = []
global branch
glob = self.global_avgpool(x)
global_triplet_feature = self.global_reduction(glob).squeeze()
global_softmax_class = self.global_softmax(global_triplet_feature)
softmax_features.append(global_softmax_class)
triplet_features.append(global_triplet_feature)
predict.append(global_triplet_feature)
# assert x.size(2) % 6 == 0
# stripes_h = int(x.size(2) / 6)
#
# for i in range(6):
# feat = F.avg_pool2d(x[:, :, i * stripes_h:(i + 1) * stripes_h, :], (stripes_h, x.size(-1)))
# feat = self.global_conv_list[i](feat)
# feat =feat.view(feat.size(0), -1)
# triplet_features.append(feat)
#
# logit = self.global_softmax_list[i](feat)
# softmax_features.append(logit)
#
#
# tmp = torch.stack(triplet_features)
# tmp = tmp.permute(1, 0, 2).contiguous()
#
# tmp = tmp.view(tmp.size(0), -1)
# predict.append(tmp)
# part branch
x = self.res_part2(x)
x = self.batch_crop(x)
triplet_feature = self.part_maxpool(x).squeeze()
feature = self.reduction(triplet_feature)
softmax_feature = self.softmax(feature)
triplet_features.append(feature)
softmax_features.append(softmax_feature)
predict.append(feature)
if self.training:
return triplet_features, softmax_features
else:
return torch.cat(predict, 1)
def get_optim_policy(self):
params = [
{
'params': self.backbone.parameters()
},
{
'params': self.res_part.parameters()
},
{
'params': self.global_reduction.parameters()
},
{
'params': self.global_softmax.parameters()
},
{
'params': self.res_part2.parameters()
},
{
'params': self.reduction.parameters()
},
{
'params': self.softmax.parameters()
},
]
return params