def __init__(self,netPath = cfg.MODEL_PATH):
super().__init__()
self.net = NET.MainNet(cfg.CLASS_NUM).to(cfg.device)
self.net.load_state_dict(torch.load(netPath))
self.net.eval()
loss_obj_cls = cls_loss_fn(output_obj[:, 5:], target_obj[:, 5].long())
loss_obj = loss_obj_conf + loss_obj_center + loss_obj_wh + loss_obj_cls
#负样本
los_noobj = con_loss_fn(output_noobj[:, 4], target_noobj[:, 4])
return alpaha * loss_obj + (1 - alpaha) * los_noobj
if __name__ == '__main__':
data_sets = dateset.MyDataset()
train_loader = DataLoader(data_sets, batch_size=4, shuffle=True)
writer = SummaryWriter()
net = NET.MainNet(cfg.CLASS_NUM).to(cfg.device)
if os.path.exists(cfg.MODEL_PATH):
print("netLoad")
net.load_state_dict(torch.load(cfg.MODEL_PATH))
# for name, value in net.named_parameters():
# if name[:8] == "trunk_13":
# print(value)
opt = torch.optim.Adam(net.parameters())
loss = None
for epoch in range(100):
for target_13, target_26, target_52, imgData in train_loader:
ouput_13, ouput_26, ouput_52 = net(imgData)