def __init__(self, config):
super(tiny_yolo, self).__init__()
self.config = config
self.loss_names = [
"x", "y", "w", "h", "conf", "cls", "recall", "precision"
]
self.seen = 0
self.header_info = np.array([0, 0, 0, self.seen, 0])
self.conv_bn = [0, 4, 8, 12, 16, 20, 24, 27, 30, 36, 41]
self.conv = [33, 44]
self.cnn = nn.Sequential(
OrderedDict([
# 0 conv 0-2
('conv0', nn.Conv2d(3, 16, 3, 1, 1, bias=False)),
('bn0', nn.BatchNorm2d(16)),
('leaky0', nn.LeakyReLU(0.1, inplace=True)),
# 1 max 3
('max1', nn.MaxPool2d(2, 2)),
# 2 conv 4-6
('conv2', nn.Conv2d(16, 32, 3, 1, 1, bias=False)),
('bn2', nn.BatchNorm2d(32)),
('leaky2', nn.LeakyReLU(0.1, inplace=True)),
# 3 max 7
('pool3', nn.MaxPool2d(2, 2)),
# 4 conv 8-10
('conv4', nn.Conv2d(32, 64, 3, 1, 1, bias=False)),
('bn4', nn.BatchNorm2d(64)),
('leaky4', nn.LeakyReLU(0.1, inplace=True)),
# 5 max 11
('pool5', nn.MaxPool2d(2, 2)),
# 6 conv 12-14
('conv6', nn.Conv2d(64, 128, 3, 1, 1, bias=False)),
('bn6', nn.BatchNorm2d(128)),
('leaky6', nn.LeakyReLU(0.1, inplace=True)),
# 7 max 15
('pool7', nn.MaxPool2d(2, 2)),
# 8 conv 16-18
('conv8', nn.Conv2d(128, 256, 3, 1, 1, bias=False)),
('bn8', nn.BatchNorm2d(256)),
('leaky8', nn.LeakyReLU(0.1, inplace=True)),
# 9 max 19
('pool9', nn.MaxPool2d(2, 2)),
# 10 conv 20-22
('conv10', nn.Conv2d(256, 512, 3, 1, 1, bias=False)),
('bn10', nn.BatchNorm2d(512)),
('leaky10', nn.LeakyReLU(0.1, inplace=True)),
# 11 max 23
('pool11', MaxPoolStride1()),
# 12 conv 24-26
('conv12', nn.Conv2d(512, 1024, 3, 1, 1, bias=False)),
('bn12', nn.BatchNorm2d(1024)),
('leaky12', nn.LeakyReLU(0.1, inplace=True)),
# 13 conv 27-29
('conv13', nn.Conv2d(1024, 256, 1, 1, 0, bias=False)),
('bn13', nn.BatchNorm2d(256)),
('leaky13', nn.LeakyReLU(0.1, inplace=True)),
# 14 conv 30-32
('conv14', nn.Conv2d(256, 512, 3, 1, 1, bias=False)),
('bn14', nn.BatchNorm2d(512)),
('leaky14', nn.LeakyReLU(0.1, inplace=True)),
# 15 conv 33
('conv15',
nn.Conv2d(512, 255, kernel_size=1, stride=1, padding=0)),
# 16 yolo 34
('yolo16', YoloLayer([3, 4, 5], self.config)),
# 17 route 35
('route17', EmptyModule()),
# 18 conv 36-38
('conv18',
nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0)),
('bn18', nn.BatchNorm2d(128)),
('leaky18', nn.LeakyReLU(0.1, inplace=True)),
# 19 upsample 39
('upsample', nn.Upsample(scale_factor=2)),
# 20 route 40
('route20', EmptyModule()),
# 21 conv 41-43
('conv21', nn.Conv2d(384, 256, 3, 1, 1, bias=False)),
('bn21', nn.BatchNorm2d(256)),
('leaky21', nn.LeakyReLU(0.1, inplace=True)),
# 22 conv 44
('conv22',
nn.Conv2d(256, 255, kernel_size=1, stride=1, padding=0)),
# 23 yolo 45
('yolo23', YoloLayer([0, 1, 2], self.config)),
]))
def create_network(self, blocks):
models = nn.ModuleList()
prev_filters = 3
out_filters = []
prev_stride = 1
out_strides = []
conv_id = 0
for block in blocks:
if block['type'] == 'net':
prev_filters = int(block['channels'])
continue
elif block['type'] == 'convolutional':
conv_id = conv_id + 1
batch_normalize = int(block['batch_normalize'])
filters = int(block['filters'])
kernel_size = int(block['size'])
stride = int(block['stride'])
is_pad = int(block['pad'])
pad = (kernel_size - 1) / 2 if is_pad else 0
activation = block['activation']
model = nn.Sequential()
if batch_normalize:
model.add_module(
'conv{0}'.format(conv_id),
nn.Conv2d(prev_filters,
filters,
kernel_size,
stride,
pad,
bias=False))
model.add_module('bn{0}'.format(conv_id),
nn.BatchNorm2d(filters))
#model.add_module('bn{0}'.format(conv_id), BN2d(filters))
else:
model.add_module(
'conv{0}'.format(conv_id),
nn.Conv2d(prev_filters, filters, kernel_size, stride,
pad))
if activation == 'leaky':
model.add_module('leaky{0}'.format(conv_id),
nn.LeakyReLU(0.1, inplace=True))
elif activation == 'relu':
model.add_module('relu{0}'.format(conv_id),
nn.ReLU(inplace=True))
prev_filters = filters
out_filters.append(prev_filters)
prev_stride = stride * prev_stride
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'maxpool':
pool_size = int(block['size'])
stride = int(block['stride'])
if stride > 1:
model = nn.MaxPool2d(pool_size, stride)
else:
model = MaxPoolStride1()
out_filters.append(prev_filters)
prev_stride = stride * prev_stride
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'avgpool':
model = GlobalAvgPool2d()
out_filters.append(prev_filters)
models.append(model)
elif block['type'] == 'softmax':
model = nn.Softmax()
out_strides.append(prev_stride)
out_filters.append(prev_filters)
models.append(model)
elif block['type'] == 'cost':
if block['_type'] == 'sse':
model = nn.MSELoss(size_average=True)
elif block['_type'] == 'L1':
model = nn.L1Loss(size_average=True)
elif block['_type'] == 'smooth':
model = nn.SmoothL1Loss(size_average=True)
out_filters.append(1)
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'reorg':
stride = int(block['stride'])
prev_filters = stride * stride * prev_filters
out_filters.append(prev_filters)
prev_stride = prev_stride * stride
out_strides.append(prev_stride)
models.append(Reorg(stride))
elif block['type'] == 'upsample':
stride = int(block['stride'])
out_filters.append(prev_filters)
prev_stride = prev_stride / stride
out_strides.append(prev_stride)
#models.append(nn.Upsample(scale_factor=stride, mode='nearest'))
models.append(Upsample(stride))
elif block['type'] == 'route':
layers = block['layers'].split(',')
ind = len(models)
layers = [
int(i) if int(i) > 0 else int(i) + ind for i in layers
]
if len(layers) == 1:
prev_filters = out_filters[layers[0]]
prev_stride = out_strides[layers[0]]
elif len(layers) == 2:
assert (layers[0] == ind - 1)
prev_filters = out_filters[layers[0]] + out_filters[
layers[1]]
prev_stride = out_strides[layers[0]]
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(EmptyModule())
elif block['type'] == 'shortcut':
ind = len(models)
prev_filters = out_filters[ind - 1]
out_filters.append(prev_filters)
prev_stride = out_strides[ind - 1]
out_strides.append(prev_stride)
models.append(EmptyModule())
elif block['type'] == 'connected':
filters = int(block['output'])
if block['activation'] == 'linear':
model = nn.Linear(prev_filters, filters)
elif block['activation'] == 'leaky':
model = nn.Sequential(nn.Linear(prev_filters, filters),
nn.LeakyReLU(0.1, inplace=True))
elif block['activation'] == 'relu':
model = nn.Sequential(nn.Linear(prev_filters, filters),
nn.ReLU(inplace=True))
prev_filters = filters
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'region':
loss = RegionLoss()
anchors = block['anchors'].split(',')
loss.anchors = [float(i) for i in anchors]
loss.num_classes = int(block['classes'])
loss.num_anchors = int(block['num'])
loss.anchor_step = len(loss.anchors) / loss.num_anchors
loss.object_scale = float(block['object_scale'])
loss.noobject_scale = float(block['noobject_scale'])
loss.class_scale = float(block['class_scale'])
loss.coord_scale = float(block['coord_scale'])
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(loss)
elif block['type'] == 'yolo':
yolo_layer = YoloLayer()
anchors = block['anchors'].split(',')
anchor_mask = block['mask'].split(',')
yolo_layer.anchor_mask = [int(i) for i in anchor_mask]
yolo_layer.anchors = [float(i) for i in anchors]
yolo_layer.num_classes = int(block['classes'])
yolo_layer.num_anchors = int(block['num'])
yolo_layer.anchor_step = len(
yolo_layer.anchors) / yolo_layer.num_anchors
yolo_layer.stride = prev_stride
#yolo_layer.object_scale = float(block['object_scale'])
#yolo_layer.noobject_scale = float(block['noobject_scale'])
#yolo_layer.class_scale = float(block['class_scale'])
#yolo_layer.coord_scale = float(block['coord_scale'])
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(yolo_layer)
else:
print('unknown type %s' % (block['type']))
return models
def create_network(self, blocks):
models = nn.ModuleList()
prev_filters = 3
out_filters = []
prev_stride = 1
out_strides = []
conv_id = 0
ind = -2
for block in blocks:
ind += 1
if block['type'] == 'net':
prev_filters = int(block['channels'])
self.width = int(block['width'])
self.height = int(block['height'])
continue
elif block['type'] == 'densenet':
class FeatureDenseNet121(nn.Module):
def __init__(self):
super(FeatureDenseNet121, self).__init__()
self.densenet121 = torchvision.models.densenet121()
def forward(self, x):
x = self.densenet121.features(x)
return x
def load_my_state_dict(net, state_dict):
own_state = net.state_dict()
for name, param in state_dict.items():
if name not in own_state:
continue
if isinstance(param, nn.Parameter):
# backwards compatibility for serialized parameters
param = param.data
own_state[name].copy_(param)
return net
def dfs_freeze(model):
for name, child in model.named_children():
for param in child.parameters():
param.requires_grad = False
dfs_freeze(child)
def dfs_unfreeze(model):
for name, child in model.named_children():
for param in child.parameters():
param.requires_grad = True
dfs_unfreeze(child)
model = FeatureDenseNet121()
dfs_unfreeze(model)
self.pretrained_path = block['pretrained_path']
gpu = None if torch.cuda.is_available() else 'cpu'
checkpoint = torch.load(self.pretrained_path, map_location=gpu)
load_my_state_dict(model, checkpoint)
print("=> loaded chexnet weights")
prev_filters = int(block['filters'])
stride = int(block['stride'])
out_filters.append(prev_filters)
prev_stride = stride * prev_stride
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'convolutional':
conv_id = conv_id + 1
batch_normalize = int(block['batch_normalize'])
filters = int(block['filters'])
kernel_size = int(block['size'])
stride = int(block['stride'])
is_pad = int(block['pad'])
pad = (kernel_size - 1) // 2 if is_pad else 0
activation = block['activation']
model = nn.Sequential()
if batch_normalize:
model.add_module(
'conv{0}'.format(conv_id),
nn.Conv2d(prev_filters,
filters,
kernel_size,
stride,
pad,
bias=False))
model.add_module('bn{0}'.format(conv_id),
nn.BatchNorm2d(filters))
#model.add_module('bn{0}'.format(conv_id), BN2d(filters))
else:
model.add_module(
'conv{0}'.format(conv_id),
nn.Conv2d(prev_filters, filters, kernel_size, stride,
pad))
if activation == 'leaky':
model.add_module('leaky{0}'.format(conv_id),
nn.LeakyReLU(0.1, inplace=True))
elif activation == 'relu':
model.add_module('relu{0}'.format(conv_id),
nn.ReLU(inplace=True))
prev_filters = filters
out_filters.append(prev_filters)
prev_stride = stride * prev_stride
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'maxpool':
pool_size = int(block['size'])
stride = int(block['stride'])
if stride > 1:
model = nn.MaxPool2d(pool_size, stride)
else:
model = MaxPoolStride1()
out_filters.append(prev_filters)
prev_stride = stride * prev_stride
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'avgpool':
model = GlobalAvgPool2d()
out_filters.append(prev_filters)
models.append(model)
elif block['type'] == 'softmax':
model = nn.Softmax()
out_strides.append(prev_stride)
out_filters.append(prev_filters)
models.append(model)
elif block['type'] == 'cost':
if block['_type'] == 'sse':
model = nn.MSELoss(size_average=True)
elif block['_type'] == 'L1':
model = nn.L1Loss(size_average=True)
elif block['_type'] == 'smooth':
model = nn.SmoothL1Loss(size_average=True)
out_filters.append(1)
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'reorg':
stride = int(block['stride'])
prev_filters = stride * stride * prev_filters
out_filters.append(prev_filters)
prev_stride = prev_stride * stride
out_strides.append(prev_stride)
models.append(Reorg(stride))
elif block['type'] == 'upsample':
stride = int(block['stride'])
out_filters.append(prev_filters)
prev_stride = prev_stride / stride
out_strides.append(prev_stride)
#models.append(nn.Upsample(scale_factor=stride, mode='nearest'))
models.append(Upsample(stride))
elif block['type'] == 'route':
layers = block['layers'].split(',')
ind = len(models)
layers = [
int(i) if int(i) > 0 else int(i) + ind for i in layers
]
if len(layers) == 1:
prev_filters = out_filters[layers[0]]
prev_stride = out_strides[layers[0]]
elif len(layers) == 2:
assert (layers[0] == ind - 1)
prev_filters = out_filters[layers[0]] + out_filters[
layers[1]]
prev_stride = out_strides[layers[0]]
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(EmptyModule())
elif block['type'] == 'shortcut':
ind = len(models)
prev_filters = out_filters[ind - 1]
out_filters.append(prev_filters)
prev_stride = out_strides[ind - 1]
out_strides.append(prev_stride)
models.append(EmptyModule())
elif block['type'] == 'connected':
filters = int(block['output'])
if block['activation'] == 'linear':
model = nn.Linear(prev_filters, filters)
elif block['activation'] == 'leaky':
model = nn.Sequential(nn.Linear(prev_filters, filters),
nn.LeakyReLU(0.1, inplace=True))
elif block['activation'] == 'relu':
model = nn.Sequential(nn.Linear(prev_filters, filters),
nn.ReLU(inplace=True))
prev_filters = filters
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'region':
region_layer = RegionLayer(use_cuda=self.use_cuda)
anchors = block['anchors'].split(',')
region_layer.anchors = [float(i) for i in anchors]
region_layer.num_classes = int(block['classes'])
region_layer.num_anchors = int(block['num'])
region_layer.anchor_step = len(
region_layer.anchors) // region_layer.num_anchors
region_layer.rescore = int(block['rescore'])
region_layer.object_scale = float(block['object_scale'])
region_layer.noobject_scale = float(block['noobject_scale'])
region_layer.class_scale = float(block['class_scale'])
region_layer.coord_scale = float(block['coord_scale'])
region_layer.thresh = float(block['thresh'])
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(region_layer)
elif block['type'] == 'yolo':
yolo_layer = YoloLayer(use_cuda=self.use_cuda)
anchors = block['anchors'].split(',')
anchor_mask = block['mask'].split(',')
yolo_layer.anchor_mask = [int(i) for i in anchor_mask]
yolo_layer.anchors = [float(i) for i in anchors]
yolo_layer.num_classes = int(block['classes'])
yolo_layer.num_anchors = int(block['num'])
yolo_layer.anchor_step = len(
yolo_layer.anchors) // yolo_layer.num_anchors
try:
yolo_layer.rescore = int(block['rescore'])
except:
pass
yolo_layer.ignore_thresh = float(block['ignore_thresh'])
yolo_layer.truth_thresh = float(block['truth_thresh'])
yolo_layer.stride = prev_stride
yolo_layer.nth_layer = ind
yolo_layer.net_width = self.width
yolo_layer.net_height = self.height
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(yolo_layer)
else:
print('unknown type %s' % (block['type']))
return models
def valid(epoch_lis, classes, draw=True, cuda=True, anchors=[]):
writer = SummaryWriter(log_dir='valid_logs',flush_secs=60)
epoch_size_val = num_val // gpu_batch
model = YoloBody(len(anchors[0]), num_classes)
anchor_masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
yolo_decodes = []
anchors = anchors.reshape([-1])
for i in range(3):
head = YoloLayer((Cfg.width, Cfg.height), anchor_masks, len(classes),
anchors, anchors.shape[0] // 2).eval()
yolo_decodes.append(head)
if Use_Data_Loader:
val_dataset = TestDataset(lines[num_train:], (input_shape[0], input_shape[1]))
gen_val = DataLoader(val_dataset, batch_size=gpu_batch, num_workers=8, pin_memory=True,
drop_last=True, collate_fn=test_dataset_collate)
else:
gen_val = TestGenerator(gpu_batch, lines[num_train:],
(input_shape[0], input_shape[1])).generate()
for epo in epoch_lis:
pth_path = find_pth_by_epoch(epo)
if not pth_path:
print('pth_path is error...')
return False
model = load_model_pth(model, pth_path)
cudnn.benchmark = True
model = model.cuda()
model.eval()
with tqdm(total=epoch_size_val, mininterval=0.3) as pbar:
infos = []
for i, batch in enumerate(gen_val):
images_src, images, targets, shapes = batch[0], batch[1], batch[2], batch[3]
with torch.no_grad():
if cuda:
images_val = Variable(torch.from_numpy(images).type(torch.FloatTensor)).cuda()
else:
images_val = Variable(torch.from_numpy(images).type(torch.FloatTensor))
outputs = model(images_val)
output_list = []
for i in range(3):
output_list.append(yolo_decodes[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output, len(classes),
conf_thres=Cfg.confidence,
nms_thres=Cfg.nms_thresh)
boxs = [box.cpu().numpy() for box in batch_detections]
# boxs = utils.post_processing(images_val, Cfg.confidence, Cfg.nms_thresh, outputs)
infos.append([boxs, targets, shapes])
if draw:
for x in range(len(boxs)):
os.makedirs('result_%d'%epo, exist_ok=True)
savename = os.path.join('result_%d'%epo, os.path.split(shapes[x][2])[-1])
plot_boxes_cv2(images_src[x], boxs[x], savename=savename, class_names=class_names)
pbar.update(1)
print()
print('===========================================================================================================')
print('++++++++cur valid epoch %d, pth_name: %s++++++++'%(epo, pth_path))
Map = make_labels_and_compute_map(infos, classes, Cfg.input_dir, save_err_miss=Cfg.save_err_mis)
writer.add_scalar('MAP/epoch', Map, epo)
print()
return True
def __init__(self, output_ch, n_classes, inference=False):
super().__init__()
self.inference = inference
self.conv1 = Conv_Bn_Activation(128, 256, 3, 1, 'leaky')
self.conv2 = Conv_Bn_Activation(256,
output_ch,
1,
1,
'linear',
bn=False,
bias=True)
self.yolo1 = YoloLayer(anchor_mask=[0, 1, 2],
num_classes=n_classes,
anchors=[
12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72,
146, 142, 110, 192, 243, 459, 401
],
num_anchors=9,
stride=8)
# R -4
self.conv3 = Conv_Bn_Activation(128, 256, 3, 2, 'leaky')
# R -1 -16
self.conv4 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
self.conv5 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
self.conv6 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
self.conv7 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
self.conv8 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
self.conv9 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
self.conv10 = Conv_Bn_Activation(512,
output_ch,
1,
1,
'linear',
bn=False,
bias=True)
self.yolo2 = YoloLayer(anchor_mask=[3, 4, 5],
num_classes=n_classes,
anchors=[
12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72,
146, 142, 110, 192, 243, 459, 401
],
num_anchors=9,
stride=16)
# R -4
self.conv11 = Conv_Bn_Activation(256, 512, 3, 2, 'leaky')
# R -1 -37
self.conv12 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
self.conv13 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
self.conv14 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
self.conv15 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
self.conv16 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
self.conv17 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
self.conv18 = Conv_Bn_Activation(1024,
output_ch,
1,
1,
'linear',
bn=False,
bias=True)
self.yolo3 = YoloLayer(anchor_mask=[6, 7, 8],
num_classes=n_classes,
anchors=[
12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72,
146, 142, 110, 192, 243, 459, 401
],
num_anchors=9,
stride=32)
def create_network(self, blocks):
models = nn.ModuleList()
for i, block in enumerate(blocks):
if block['type'] == 'net':
self.width = int(block['width'])
self.height = int(block['height'])
continue
elif block['type'] == 'yolo':
yololayer = YoloLayer(use_cuda=self.use_cuda)
anchors = block['anchors'].split(',')
anchor_mask = block['mask'].split(',')
yololayer.anchor_mask = [int(i) for i in anchor_mask]
yololayer.anchors = [float(i) for i in anchors]
yololayer.num_classes = int(block['classes'])
yololayer.num_anchors = int(block['num'])
yololayer.anchor_step = len(yololayer.anchors)//yololayer.num_anchors
try:
yololayer.rescore = int(block['rescore'])
except:
pass
yololayer.nth_layer = i
yololayer.ignore_thresh = float(block['ignore_thresh'])
yololayer.truth_thresh = float(block['truth_thresh'])
yololayer.net_width = self.width
yololayer.net_height = self.height
models.append(yololayer)
return models
def create_network(self, blocks):
models = nn.ModuleList()
prev_filters = 3
out_filters = []
prev_stride = 1
out_strides = []
conv_id = 0
for block in blocks:
if block["type"] == "net":
prev_filters = int(block["channels"])
continue
elif block["type"] == "convolutional":
conv_id = conv_id + 1
batch_normalize = int(block["batch_normalize"])
filters = int(block["filters"])
kernel_size = int(block["size"])
stride = int(block["stride"])
is_pad = int(block["pad"])
pad = (kernel_size - 1) // 2 if is_pad else 0
activation = block["activation"]
model = nn.Sequential()
if batch_normalize:
model.add_module(
"conv{0}".format(conv_id),
nn.Conv2d(prev_filters,
filters,
kernel_size,
stride,
pad,
bias=False))
model.add_module("bn{0}".format(conv_id),
nn.BatchNorm2d(filters))
#model.add_module("bn{0}".format(conv_id), BN2d(filters))
else:
model.add_module(
"conv{0}".format(conv_id),
nn.Conv2d(prev_filters, filters, kernel_size, stride,
pad))
if activation == "leaky":
model.add_module("leaky{0}".format(conv_id),
nn.LeakyReLU(0.1, inplace=True))
elif activation == "relu":
model.add_module("relu{0}".format(conv_id),
nn.ReLU(inplace=True))
prev_filters = filters
out_filters.append(prev_filters)
prev_stride = stride * prev_stride
out_strides.append(prev_stride)
models.append(model)
elif block["type"] == "maxpool":
pool_size = int(block["size"])
stride = int(block["stride"])
if stride > 1:
model = nn.MaxPool2d(pool_size, stride)
else:
model = MaxPoolStride1()
out_filters.append(prev_filters)
prev_stride = stride * prev_stride
out_strides.append(prev_stride)
models.append(model)
elif block["type"] == "avgpool":
model = GlobalAvgPool2d()
out_filters.append(prev_filters)
models.append(model)
elif block["type"] == "softmax":
model = nn.Softmax()
out_strides.append(prev_stride)
out_filters.append(prev_filters)
models.append(model)
elif block["type"] == "cost":
if block["_type"] == "sse":
model = nn.MSELoss(size_average=True)
elif block["_type"] == "L1":
model = nn.L1Loss(size_average=True)
elif block["_type"] == "smooth":
model = nn.SmoothL1Loss(size_average=True)
out_filters.append(1)
out_strides.append(prev_stride)
models.append(model)
elif block["type"] == "reorg":
stride = int(block["stride"])
prev_filters = stride * stride * prev_filters
out_filters.append(prev_filters)
prev_stride = prev_stride * stride
out_strides.append(prev_stride)
models.append(Reorg(stride))
elif block["type"] == "upsample":
stride = int(block["stride"])
out_filters.append(prev_filters)
prev_stride = prev_stride // stride
out_strides.append(prev_stride)
#models.append(nn.Upsample(scale_factor=stride, mode="nearest"))
models.append(Upsample(stride))
elif block["type"] == "route":
layers = block["layers"].split(",")
ind = len(models)
layers = [
int(i) if int(i) > 0 else int(i) + ind for i in layers
]
if len(layers) == 1:
prev_filters = out_filters[layers[0]]
prev_stride = out_strides[layers[0]]
elif len(layers) == 2:
assert (layers[0] == ind - 1)
prev_filters = out_filters[layers[0]] + out_filters[
layers[1]]
prev_stride = out_strides[layers[0]]
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(EmptyModule())
elif block["type"] == "shortcut":
ind = len(models)
prev_filters = out_filters[ind - 1]
out_filters.append(prev_filters)
prev_stride = out_strides[ind - 1]
out_strides.append(prev_stride)
models.append(EmptyModule())
elif block["type"] == "connected":
filters = int(block["output"])
if block["activation"] == "linear":
model = nn.Linear(prev_filters, filters)
elif block["activation"] == "leaky":
model = nn.Sequential(nn.Linear(prev_filters, filters),
nn.LeakyReLU(0.1, inplace=True))
elif block["activation"] == "relu":
model = nn.Sequential(nn.Linear(prev_filters, filters),
nn.ReLU(inplace=True))
prev_filters = filters
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(model)
elif block["type"] == "region":
loss = RegionLoss()
anchors = block["anchors"].split(",")
loss.anchors = [float(i) for i in anchors]
loss.num_classes = int(block["classes"])
loss.num_anchors = int(block["num"])
loss.anchor_step = len(loss.anchors) // loss.num_anchors
loss.object_scale = float(block["object_scale"])
loss.noobject_scale = float(block["noobject_scale"])
loss.class_scale = float(block["class_scale"])
loss.coord_scale = float(block["coord_scale"])
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(loss)
elif block["type"] == "yolo":
yolo_layer = YoloLayer()
anchors = block["anchors"].split(",")
anchor_mask = block["mask"].split(",")
yolo_layer.anchor_mask = [int(i) for i in anchor_mask]
yolo_layer.anchors = [float(i) for i in anchors]
yolo_layer.num_classes = int(block["classes"])
yolo_layer.num_anchors = int(block["num"])
yolo_layer.anchor_step = len(
yolo_layer.anchors) // yolo_layer.num_anchors
yolo_layer.stride = prev_stride
#yolo_layer.object_scale = float(block["object_scale"])
#yolo_layer.noobject_scale = float(block["noobject_scale"])
#yolo_layer.class_scale = float(block["class_scale"])
#yolo_layer.coord_scale = float(block["coord_scale"])
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(yolo_layer)
else:
print("unknown type %s" % (block["type"]))
return models