def main():
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
model = efficientdet_d0(False, num_classes=len(labels_map))
optim = torch.optim.SGD(model.parameters(), 0, 0.9, weight_decay=4e-5)
# 数据集自行合并(The dataset merges by yourself)
# 链接：https://pan.baidu.com/s/17iop7UBnSGExW64cip-pYw
# 提取码：sdvx
train_dataset = VOC_Dataset(voc_dir,
"0712",
"trainval",
transforms=train_transforms)
test_dataset = VOC_Dataset(voc_dir,
"0712",
"test",
transforms=test_transforms)
ap_counter = APCounter(labels_map, 0.5)
writer = SummaryWriter(comment=comment)
logger = Logger(50, writer)
checker = Checker(
Tester(model, train_dataset, batch_size, device, ap_counter, 1000),
Tester(model, test_dataset, batch_size, device, ap_counter, 5000),
Saver(model), logger, 8, 1)
lr_scheduler = LRScheduler(optim, lr_func)
trainer = Trainer(model, optim, train_dataset, batch_size, device,
lr_scheduler, logger, checker)
print("配置: %s" % comment, flush=True)
trainer.train((0, 40))
writer.close()
# read images(假定)
image_fname = "images/1.png"
with Image.open(image_fname) as image:
image_o = pil_to_cv(image)
image_1 = trans.ToTensor()(image).to(device) # 图片1
del image
image_2 = torch.rand(3, 600, 800).to(device) # 图片2
# target(假定)
targets = [{
"labels":
torch.tensor([0, 1]).to(device),
"boxes":
torch.tensor([[0., 10., 200., 200.], [100., 20., 500., 100.]]).to(device)
}, {
"labels": torch.tensor([]).to(device).reshape((0, )),
"boxes": torch.tensor([]).to(device).reshape((0, 4))
}]
# train
model = efficientdet_d0(True, norm_layer=FrozenBatchNorm2d).to(device)
optim = torch.optim.Adam(model.parameters(), 5e-4)
for i in range(20):
optim.zero_grad()
loss = model([image_1, image_2], targets)
loss_sum = sum(loss.values())
loss_sum.backward()
optim.step()
print("loss: %s" % loss)
# target(假定assumption)
# 由于是随机标签、y_pred 会反复横跳, 此文件只是展示输入输出格式
targets = [
{
"labels":
torch.tensor([0, 1]).to(device),
"boxes":
torch.tensor([[50., 50., 800., 800.], [200., 200., 1300.,
1300.]]).to(device)
},
{
"labels": torch.tensor([]).to(device).reshape(
(0, )), # 支持空标签图片(Support empty label images)
"boxes": torch.tensor([]).to(device).reshape((0, 4))
}
]
# train
model = efficientdet_d0(True).to(device)
optim = torch.optim.SGD(model.parameters(),
2e-3,
0.9,
weight_decay=1e-4,
nesterov=True)
for i in range(10):
optim.zero_grad()
loss = model([image_1, image_2], targets)
loss_sum = sum(loss.values())
loss_sum.backward()
optim.step()
print("loss: %s" % loss)
# Author: Jintao Huang
# Time: 2020-5-21
from models.efficientdet import efficientdet_d0
import torch
from utils.detection.predictor import Predictor
# -------------------------------- 参数
image_path = "images/1.png"
score_thresh = 0.2
nms_thresh = 0.2
# --------------------------------
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
# pred
model = efficientdet_d0(True)
predictor = Predictor(model, device, None, None)
predictor.pred_image_and_save(image_path, None, "max", score_thresh, nms_thresh)
# predictor.pred_image_and_show(image_path, "max", score_thresh, nms_thresh) # test function