def __init__(self, weight_path=None, resume=False):
super(Build_Model, self).__init__()
self.__anchors = torch.FloatTensor(cfg.MODEL["ANCHORS"])
self.__strides = torch.FloatTensor(cfg.MODEL["STRIDES"])
if cfg.TRAIN["DATA_TYPE"] == 'VOC':
self.__nC = cfg.VOC_DATA["NUM"]
elif cfg.TRAIN["DATA_TYPE"] == 'COCO':
self.__nC = cfg.COCO_DATA["NUM"]
else:
self.__nC = cfg.Customer_DATA["NUM"]
self.__out_channel = cfg.MODEL["ANCHORS_PER_SCLAE"] * (self.__nC + 5)
self.__yolov4 = YOLOv4(weight_path=weight_path,
out_channels=self.__out_channel,
resume=resume)
# small
self.__head_s = Yolo_head(nC=self.__nC,
anchors=self.__anchors[0],
stride=self.__strides[0])
# medium
self.__head_m = Yolo_head(nC=self.__nC,
anchors=self.__anchors[1],
stride=self.__strides[1])
# large
self.__head_l = Yolo_head(nC=self.__nC,
anchors=self.__anchors[2],
stride=self.__strides[2])
def __init__(self, init_weights=True):
super(Build_Model, self).__init__()
self.__anchors = torch.FloatTensor(cfg.MODEL["ANCHORS"])
self.__strides = torch.FloatTensor(cfg.MODEL["STRIDES"])
self.__nC = cfg.VOC_DATA["NUM"] if cfg.TRAIN[
"DATA_TYPE"] == 'VOC' else cfg.COCO_DATA["NUM"]
self.__out_channel = cfg.MODEL["ANCHORS_PER_SCLAE"] * (self.__nC + 5)
self.__yolov4 = YOLOv4(out_channels=self.__out_channel)
# small
self.__head_s = Yolo_head(nC=self.__nC,
anchors=self.__anchors[0],
stride=self.__strides[0])
# medium
self.__head_m = Yolo_head(nC=self.__nC,
anchors=self.__anchors[1],
stride=self.__strides[1])
# large
self.__head_l = Yolo_head(nC=self.__nC,
anchors=self.__anchors[2],
stride=self.__strides[2])
if init_weights:
self.__init_weights()
lr = self.__lr_max / self.__warmup * t
else:
T_max = self.__T_max - self.__warmup
t = t - self.__warmup
lr = self.__lr_min + 0.5 * (self.__lr_max - self.__lr_min) * (
1 + np.cos(t / T_max * np.pi))
for param_group in self.__optimizer.param_groups:
param_group["lr"] = lr
if __name__ == "__main__":
import matplotlib.pyplot as plt
from model.YOLOv4 import YOLOv4
import torch.optim as optim
net = YOLOv4()
optimizer = optim.SGD(net.parameters(), 1e-4, 0.9, weight_decay=0.0005)
scheduler = CosineDecayLR(optimizer, 50 * 2068, 1e-4, 1e-6, 2 * 2068)
# Plot lr schedule
y = []
for t in range(50):
for i in range(2068):
scheduler.step(2068 * t + i)
y.append(optimizer.param_groups[0]["lr"])
print(y)
plt.figure()
plt.plot(y, label="LambdaLR")
plt.xlabel("steps")
plt.ylabel("LR")
if tf.test.gpu_device_name():
print('Default GPU Device:{}'.format(tf.test.gpu_device_name()))
else:
print("Please install GPU version of TF")
from model.YOLOv4 import YOLOv4
import numpy as np
import gc
from tqdm import tqdm
from gpuinfo import GPUInfo
x = 320
y = 320
inp = np.ones((4, x, y, 3), dtype=np.float64)
network = YOLOv4(side=x)
network.compile(loss="mse", optimizer="sgd", metrics="acc")
small, medium, large = network(inp)
print(GPUInfo.get_info())
print((x, y), " work and give :", small.shape, medium.shape, large.shape)
network.summary()
print(network.layers[-1].output_shape)
from dataset import Dataset
dataset = Dataset(4,
network.layers[-1].output_shape,
shape=network.layers[0].input_shape[1:3])
print("input shape is :", dataset.shape)
image_path = dataset.get_dir("test")