Este artículo le enseñará a Shuai Bikanguan cómo reemplazar la red troncal del modelo de detección de tramas oblicuas de BBAVectors.
Este blog asume que Shuaibi ha podido usar BBAVectors para entrenar sus propios datos, pero no sabe cómo reemplazar su red troncal. El resnet101 utilizado en el código original se utiliza como red troncal. Esta red tiene altos requisitos de recursos durante el entrenamiento, entonces, ¿cómo reemplazarla con el resnet18 y mobilnet más livianos?
1. La red troncal se modifica al modelo resnet18
1. Modifique los parámetros de la red troncal en ctrbox_net.py
#修改前
# self.base_network = resnet.resnet101(pretrained=pretrained)
# self.dec_c2 = CombinationModule(512, 256, batch_norm=True)
# self.dec_c3 = CombinationModule(1024, 512, batch_norm=True)
# self.dec_c4 = CombinationModule(2048, 1024, batch_norm=True)
#修改后
self.base_network = resnet.resnet18(pretrained=pretrained)
self.dec_c2 = CombinationModule(128, 64, batch_norm=True)
self.dec_c3 = CombinationModule(256, 128, batch_norm=True)
self.dec_c4 = CombinationModule(512, 256, batch_norm=True)En primer lugar, debemos comprender que los 256, 512, 1024 y 2048 anteriores son en realidad los tamaños de canal de los cuatro mapas de características en la red troncal resnet101. Al analizar resnet18 de resnet.py, podemos ver que la cantidad de canales de sus cuatro mapas de características son respectivamente 64, 128, 256 y 512. Solo necesita agregar las siguientes líneas de código al final de resnet.py, y luego ejecute resnet.py Sabemos que solo la cantidad de canales es diferente, por lo que podemos modificar directamente los parámetros anteriores.
if __name__ == '__main__':
device='cpu'
input=np.ones((1,3,512,512)).astype(np.float32)
dummy_input = torch.from_numpy(input).to(device)
model=resnet18().to(device)
logit=model(dummy_input)
print(logit[-1].size()) #torch.Size([1, 512, 16, 16])
print(logit[-2].size()) #torch.Size([1, 256, 32, 32])
print(logit[-3].size()) #torch.Size([1, 128, 64, 64])
print(logit[-4].size()) #torch.Size([1, 64, 128, 128])2. Modificar los canales en ctrbox_net.py
Cuando down_ratio=4, el número de canales de salida de c2_combine es 256, que también es el número de canales de entrada de los siguientes cuatro cabezales. El código de cálculo de ctrbox_net.py es el siguiente
class CTRBOX(nn.Module):
def __init__(self, heads, pretrained, down_ratio, final_kernel, head_conv,export=False):
super(CTRBOX, self).__init__()
channels = [3, 64, 256, 512, 1024, 2048]
assert down_ratio in [2, 4, 8, 16]
self.l1 = int(np.log2(down_ratio))
El número de canales de entrada del cabezal eschannels[self.l1]=256. Pero después de cambiar a resnet18, la cantidad de canales emitidos por c2_combine se convierte en 64. En este momento, cambiará naturalmente el down_ratio en main.py a 2, de modo que canales[self.l1]=64.
Si modifica down_ratio=2, es un gran error. Debido a que el tamaño de salida final de ctrbox_net se reduce 4 veces el tamaño de entrada, es decir, la entrada es 512 x 512 y la salida de los cuatro cabezales es 128 x 128. Al calcular la pérdida, el tamaño de gt también debe ser 128 x 128. Down_ration también se usa en dataset/base.py.Si se cambia a 2, el tamaño de gt se calcula en 256x256 y se informará un error.
**Solución**
No cambiamos el down_ratio, después de completar la modificación de contenido en 1, cambia directamente el 256 en canales = [3, 64, 256, 512, 1024, 2048] en ctrbox_net.py a 64, es decir, canales = [3, 64, 64, 512, 1024, 2048] Puedes empezar a entrenar, violencia simple.
3. Resultados
2. Modifique a una red troncal ligera de mobilenetv2
El código de mobilenetv2 es el siguiente
import math
import os
import torch
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
BatchNorm2d = nn.BatchNorm2d
def conv_bn(inp, oup, stride):
return nn.Sequential(
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
BatchNorm2d(oup),
nn.ReLU6(inplace=True)
)
def conv_1x1_bn(inp, oup):
return nn.Sequential(
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
BatchNorm2d(oup),
nn.ReLU6(inplace=True)
)
class InvertedResidual(nn.Module):
def __init__(self, inp, oup, stride, expand_ratio):
super(InvertedResidual, self).__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = round(inp * expand_ratio)
self.use_res_connect = self.stride == 1 and inp == oup
if expand_ratio == 1:
self.conv = nn.Sequential(
# dw
nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
BatchNorm2d(hidden_dim),
nn.ReLU6(inplace=True),
# pw-linear
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
BatchNorm2d(oup),
)
else:
self.conv = nn.Sequential(
# pw
nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
BatchNorm2d(hidden_dim),
nn.ReLU6(inplace=True),
# dw
nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
BatchNorm2d(hidden_dim),
nn.ReLU6(inplace=True),
# pw-linear
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
BatchNorm2d(oup),
)
def forward(self, x):
if self.use_res_connect:
return x + self.conv(x)
else:
return self.conv(x)
class MobileNetV2(nn.Module):
def __init__(self, n_class=1000, input_size=224, width_mult=1.):
super(MobileNetV2, self).__init__()
block = InvertedResidual
input_channel = 32
last_channel = 1280
interverted_residual_setting = [
# t, c, n, s
[1, 16, 1, 1],
[6, 24, 2, 2],
[6, 32, 3, 2],
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1],
]
# building first layer
assert input_size % 32 == 0
input_channel = int(input_channel * width_mult)
self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel
self.features = [conv_bn(3, input_channel, 2)]
# building inverted residual blocks
for t, c, n, s in interverted_residual_setting:
output_channel = int(c * width_mult)
for i in range(n):
if i == 0:
self.features.append(block(input_channel, output_channel, s, expand_ratio=t))
else:
self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))
input_channel = output_channel
# building last several layers
self.features.append(conv_1x1_bn(input_channel, self.last_channel))
# make it nn.Sequential
self.features = nn.Sequential(*self.features)
# building classifier
self.classifier = nn.Sequential(
nn.Dropout(0.2),
nn.Linear(self.last_channel, n_class),
)
self._initialize_weights()
def forward(self, x):
x = self.features(x)
x = x.mean(3).mean(2)
x = self.classifier(x)
return x
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
n = m.weight.size(1)
m.weight.data.normal_(0, 0.01)
m.bias.data.zero_()
def load_url(url, model_dir='./model_data', map_location=None):
if not os.path.exists(model_dir):
os.makedirs(model_dir)
filename = url.split('/')[-1]
cached_file = os.path.join(model_dir, filename)
if os.path.exists(cached_file):
return torch.load(cached_file, map_location=map_location)
else:
return model_zoo.load_url(url,model_dir=model_dir)
def mobilenetv2(pretrained=False, **kwargs):
model = MobileNetV2(n_class=1000, **kwargs)
if pretrained:
model.load_state_dict(load_url('http://sceneparsing.csail.mit.edu/model/pretrained_resnet/mobilenet_v2.pth.tar'), strict=False)
return model
class Backbone(nn.Module):
def __init__(self,pretrained=False):
super(Backbone, self).__init__()
model=mobilenetv2(pretrained)
self.features = model.features[:-1]
def forward(self, x):
C1=self.features[:4](x)
C2=self.features[4:7](C1)
C3=self.features[7:14](C2)
C4=self.features[14:](C3)
return C1,C2,C3,C4
if __name__ == '__main__':
import numpy as np
device='cpu'
model=Backbone().to(device)
input=np.ones((1,3,512,512)).astype(np.float32)
dummy_input = torch.from_numpy(input).to(device)
logit=model(dummy_input)
print(logit[-1].size())
print(logit[-2].size())
print(logit[-3].size())
print(logit[-4].size())
Los tamaños de los cuatro mapas de características son
torch.Size([1, 320, 16, 16])
torch.Size([1, 96, 32, 32])
torch.Size([1, 32, 64, 64])
torch.Size([1, 24, 128, 128])En comparación con el cambio a resnet18 anterior, sabrá que los canales y cada canal deben cambiarse a lo siguiente
# channels = [3, 64, 256, 512, 1024, 2048] # 当下面采用resnet101的时候用这个
# self.base_network = resnet.resnet101(pretrained=pretrained)
# self.dec_c2 = CombinationModule(512, 256, batch_norm=True)
# self.dec_c3 = CombinationModule(1024, 512, batch_norm=True)
# self.dec_c4 = CombinationModule(2048, 1024, batch_norm=True)
# channels = [3, 64, 64, 512, 1024, 2048] # 用resnet18的时候是这个
# self.base_network = resnet.resnet18(pretrained=pretrained)
# self.dec_c2 = CombinationModule(128, 64, batch_norm=True)
# self.dec_c3 = CombinationModule(256, 128, batch_norm=True)
# self.dec_c4 = CombinationModule(512, 256, batch_norm=True)
channels = [3, 64, 24, 512, 1024, 2048] # mobilenetv2的时候是这个
self.base_network = mobilenet.Backbone(pretrained=pretrained)
self.dec_c2 = CombinationModule(32, 24, batch_norm=True)
self.dec_c3 = CombinationModule(96, 32, batch_norm=True)
self.dec_c4 = CombinationModule(320, 96, batch_norm=True)
2. Resultados
La caída de pérdidas no es tan buena como la pérdida de la red troncal con resnet18, que es una situación normal.
3. Código completo
Se puede hacer referencia al código completo : vea aquí , este es mi código BBAVectors modificado, que incluye DOTA_Devkit, rolabelimg xml a código de formato txt de 4 puntos, división de código de conjunto de datos y exportación de código de modelo onnx con decodificador, conversión de modelo tensorrt y razonamiento para el código, puedes leer mis otros blogs sobre el razonamiento tensorrt.