def train(self,save_path):
if self.net == 'pnet':
net = Pnet()
elif self.net == 'rnet':
net = Rnet()
elif self.net == 'onet':
net = Onet()
net = net.cuda()
# net.load_state_dict(torch.load(save_path))
optimizer = optim.Adam(net.parameters())
conf_loss_fun = nn.BCELoss()
off_loos_fun = nn.MSELoss()
for epoch in range(20000):
imgdata, conf, offset = DataSet.GetIter(self.dataloader)
imgdata = imgdata.cuda()
conf = conf.cuda()
offset = offset.cuda()
confidence,offset_out = net(imgdata)
#置信度的损失需要正负样本
#获得置信度小于2的掩码
conn_mask = torch.lt(conf,2)
#得到符合条件的置信度
conf_ = conf[conn_mask]
confidence_ = confidence[conn_mask]
#偏移的损失需要正样本和部分样本
#得到置信度大于0的掩码
off_mask = torch.gt(conf,0)
#得到符合条件的偏移
offset = offset[off_mask[:,0]]
offset_out = offset_out[off_mask[:,0]]
conf_loss = conf_loss_fun(confidence_,conf_)
off_loss = off_loos_fun(offset_out,offset)
loss = conf_loss + off_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(loss)
torch.save(net.state_dict(),save_path)
