def __init__(self, hparams):
super().__init__()
self.fc_num1 = 200
self.fc_num2 = 300
self.hparams = hparams
self.learning_rate = hparams.learning_rate
self.anchors = get_anchors(hparams.anchors_file)
self.anchors1 = np.reshape(self.anchors[0], [1, 2])
self.anchors0 = self.anchors[1:]
self.detection_classes = hparams.detection_classes
super(AutoNet, self).__init__()
self.efficientNet = EfficientNet.from_name('efficientnet-b2', freeze=hparams.freeze)
self.compressed = nn.Sequential(
nn.Conv2d(352, 16, 1, bias=False),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
)
self.compressed_1 = nn.Sequential(
nn.Conv2d(120, 8, 1, bias=False),
nn.BatchNorm2d(8),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
)
self.fc1_1_1 = nn.Sequential(
nn.Linear(16 * 8 * 10, self.fc_num1 * 2, bias=False),
nn.BatchNorm1d(self.fc_num1 * 2),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.fc1_1_2 = nn.Sequential(
nn.Linear(8 * 16 * 20, self.fc_num1 * 2, bias=False),
nn.BatchNorm1d(self.fc_num1 * 2),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.fc1_2_1 = nn.Sequential(
nn.Linear(16 * 8 * 10, self.fc_num2 * 2, bias=False),
nn.BatchNorm1d(self.fc_num2 * 2),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.fc1_2_2 = nn.Sequential(
nn.Linear(8 * 16 * 20, self.fc_num2 * 2, bias=False),
nn.BatchNorm1d(self.fc_num2 * 2),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.fc2_0_1 = nn.ModuleList([])
for i in range(6):
if i != 1 and i != 4:
self.fc2_0_1.append(nn.Sequential(
nn.Linear(self.fc_num1 * 2, 14 * 13 * 32, bias=False),
nn.BatchNorm1d(14 * 13 * 32),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
))
else:
self.fc2_0_1.append(nn.Sequential(
nn.Linear(self.fc_num1 * 2, 13 * 18 * 32, bias=False),
nn.BatchNorm1d(13 * 18 * 32),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
))
self.fc2_0_2 = nn.ModuleList([])
for i in range(6):
if i != 1 and i != 4:
self.fc2_0_2.append(nn.Sequential(
nn.Linear(self.fc_num1 * 2, 28 * 26 * 8, bias=False),
nn.BatchNorm1d(28 * 26 * 8),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
))
else:
self.fc2_0_2.append(nn.Sequential(
nn.Linear(self.fc_num1 * 2, 26 * 36 * 8, bias=False),
nn.BatchNorm1d(26 * 36 * 8),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
))
self.fc2_1_1 = nn.ModuleList([])
for i in range(6):
if i != 1 and i != 4:
self.fc2_1_1.append(nn.Sequential(
nn.Linear(self.fc_num2 * 2, 14 * 13 * 64, bias=False),
nn.BatchNorm1d(14 * 13 * 64),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
))
else:
self.fc2_1_1.append(nn.Sequential(
nn.Linear(self.fc_num2 * 2, 13 * 18 * 64, bias=False),
nn.BatchNorm1d(13 * 18 * 64),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
))
self.fc2_1_2 = nn.ModuleList([])
for i in range(6):
if i != 1 and i != 4:
self.fc2_1_2.append(nn.Sequential(
nn.Linear(self.fc_num2 * 2, 28 * 26 * 8, bias=False),
nn.BatchNorm1d(28 * 26 * 8),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
))
else:
self.fc2_1_2.append(nn.Sequential(
nn.Linear(self.fc_num2 * 2, 26 * 36 * 8, bias=False),
nn.BatchNorm1d(26 * 36 * 8),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
))
self.inplanes = 32
self.conv0 = self._make_layer(BasicBlock, 32, 2)
self.deconv0 = self._make_deconv_layer(32, 8)
self.inplanes = 16
self.conv1 = self._make_layer(BasicBlock, 16, 2)
self.deconv1 = self._make_deconv_layer(16, 8)
self.inplanes = 8
self.conv2 = self._make_layer(BasicBlock, 8, 2)
self.deconv2 = self._make_deconv_layer(8, 4)
self.inplanes = 4
self.conv3 = self._make_layer(BasicBlock, 4, 2)
self.deconv3 = self._make_deconv_layer(4, 2)
self.convfinal = nn.Conv2d(2, 2, 1)
self.inplanes = 64
self.conv0_1_detect = self._make_layer(BasicBlock, 64, 2)
self.convfinal_0 = nn.Conv2d(64, len(self.anchors0) * (self.detection_classes + 5), 1)
self.yolo0 = YOLOLayer(self.anchors0, self.detection_classes, 800)
self.conv0_1 = self._make_layer(BasicBlock, 64, 2)
self.deconv0_1 = self._make_deconv_layer(64, 8)
self.inplanes = 16
self.conv1_1_detect = self._make_layer(BasicBlock, 16, 2)
self.convfinal_1 = nn.Conv2d(16, len(self.anchors1) * (self.detection_classes + 5), 1)
self.yolo1 = YOLOLayer(self.anchors1, self.detection_classes, 800)
self.conv1_1 = self._make_layer(BasicBlock, 16, 2)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.ConvTranspose2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
def __init__(self, hparams):
super().__init__()
self.latent = 600
self.fc_num = 300
self.hparams = hparams
self.learning_rate = hparams.learning_rate
self.anchors = get_anchors(hparams.anchors_file)
self.anchors1 = np.reshape(self.anchors[0], [1, 2])
self.anchors0 = self.anchors[1:]
self.detection_classes = hparams.detection_classes
super(AutoNet, self).__init__()
self.efficientNet = EfficientNet.from_name('efficientnet-b1',
freeze=hparams.freeze)
feature = self.efficientNet._fc.in_features
self.efficientNet._fc = nn.Sequential(
nn.Linear(in_features=feature, out_features=self.latent),
nn.BatchNorm1d(self.latent),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.fc1 = nn.Sequential(
nn.Linear(self.latent, self.fc_num, bias=False),
nn.BatchNorm1d(self.fc_num),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.fc2 = nn.Sequential(
nn.Linear(self.fc_num * 6, 25 * 25 * 32, bias=False),
nn.BatchNorm1d(25 * 25 * 32),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.fc1_1 = nn.Sequential(
nn.Linear(self.latent, self.fc_num, bias=False),
nn.BatchNorm1d(self.fc_num),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.fc2_1 = nn.Sequential(
nn.Linear(self.fc_num * 6, 25 * 25 * 64, bias=False),
nn.BatchNorm1d(25 * 25 * 64),
nn.ReLU(inplace=True),
nn.Dropout(0.25),
)
self.inplanes = 32
self.conv0 = self._make_layer(BasicBlock, 32, 2)
self.deconv0 = self._make_deconv_layer(32, 16)
self.inplanes = 16
self.conv1 = self._make_layer(BasicBlock, 16, 2)
self.deconv1 = self._make_deconv_layer(16, 8)
self.inplanes = 8
self.conv2 = self._make_layer(BasicBlock, 8, 2)
self.deconv2 = self._make_deconv_layer(8, 4)
self.inplanes = 4
self.conv3 = self._make_layer(BasicBlock, 4, 2)
self.deconv3 = self._make_deconv_layer(4, 2)
self.convfinal = nn.Conv2d(2, 2, 1)
self.inplanes = 64
self.conv0_1_detect = self._make_layer(BasicBlock, 64, 2)
self.convfinal_0 = nn.Conv2d(
64,
len(self.anchors0) * (self.detection_classes + 5), 1)
self.yolo0 = YOLOLayer(self.anchors0, self.detection_classes, 800)
self.conv0_1 = self._make_layer(BasicBlock, 64, 2)
self.deconv0_1 = self._make_deconv_layer(64, 16)
self.conv0_1 = self._make_layer(BasicBlock, 64, 2)
self.inplanes = 16
self.conv1_1_detect = self._make_layer(BasicBlock, 16, 2)
self.convfinal_1 = nn.Conv2d(
16,
len(self.anchors1) * (self.detection_classes + 5), 1)
self.yolo1 = YOLOLayer(self.anchors1, self.detection_classes, 800)
self.conv1_1 = self._make_layer(BasicBlock, 16, 2)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.ConvTranspose2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
scene_batch_size=2)
trainset, testset = torch.utils.data.random_split(labeled_trainset, [
int(0.90 * len(labeled_trainset)),
len(labeled_trainset) - int(0.90 * len(labeled_trainset))
])
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=2,
shuffle=True,
num_workers=8,
collate_fn=collate_fn_lstm)
testloader = torch.utils.data.DataLoader(testset,
batch_size=2,
shuffle=True,
num_workers=8,
collate_fn=collate_fn_lstm)
anchors = get_anchors(anchor_file)
if pretrain_file is not None:
model = bothModel.trainModel(device, anchors, freeze=True)
pretrain_dict = torch.load(pretrain_file, map_location=device)
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if (k in model_dict and re.search('^efficientNet.*', k) and (
not re.search('^efficientNet._fc.*', k)))
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
else:
model = bothModel.trainModel(anchors, freeze=False, device=device)
from config.yolov4_config import TRAIN, VAL, MODEL, LR
from model.loss import yolo4_loss
from torch.optim.lr_scheduler import CosineAnnealingLR, CosineAnnealingWarmRestarts, StepLR
import torchsnooper
if __name__ == "__main__":
log_file = MODEL["LOG_PATH"] + 'logs.txt'
logger = Logger(log_file).get_log()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
epochs = TRAIN["YOLO_EPOCHS"]
batch_size = TRAIN["BATCH_SIZE"]
anchors = get_anchors(MODEL["ANCHOR_PATH"]).to(device)
strides = torch.tensor(MODEL["STRIDES"]).to(device)
path = TRAIN["PATH"]
train_path = TRAIN["TRAIN_PATH"]
val_path = VAL["VAL_PATH"]
train_img_size = TRAIN["TRAIN_IMG_SIZE"]
class_names = get_classes(MODEL["CLASS_PATH"])
num_classes = len(class_names)
pretrained = TRAIN["PRE_TRAIN"]
pretrained_weights = TRAIN["PRE_TRAIN_W"]
accumulation = TRAIN["GD_ACCUM"]
warmup_lr = TRAIN["LR_WARMUP"]
learning_rate = TRAIN["LR_INIT"]
saved_path = MODEL["WEIGHTS_SAVED_PATH"]