Ejemplo n.º 1
0
import torch
import net.checknet as nets
import tool.utils as utils
from PIL import Image
import numpy as np
import time
import cv2

if __name__ == '__main__':
    #net
    pnet = nets.P_Net(istraining=False)
    rnet = nets.R_Net(istraining=False)
    onet = nets.O_Net(istraining=False, island=False)
    if torch.cuda.is_available():
        pnet = pnet.cuda()
        rnet = rnet.cuda()
        onet = onet.cuda()
    pnet.eval()
    rnet.eval()
    onet.eval()
    pnet.load_state_dict(
        torch.load(
            r'C:\Users\liev\Desktop\myproject\face_recognition\params\p_params_ok.pkl'
        ))  # 导入训练参数
    rnet.load_state_dict(
        torch.load(
            r'C:\Users\liev\Desktop\myproject\face_recognition\params\r_params_ok.pkl'
        ))  # 导入训练参数
    onet.load_state_dict(
        torch.load(
            r'C:\Users\liev\Desktop\myproject\face_recognition\params\o_params_ok.pkl'
Ejemplo n.º 2
0
    def train(self):
        if self.net == "pnet":
            net = cnet.P_Net(istraining=True)
        elif self.net == "rnet":
            net = cnet.R_Net(istraining=True)
        elif self.net == "onet":
            net = cnet.O_Net(istraining=True, island=self.island)
        else:
            print("net is error")
            return False

        # if os.path.exists(self.parampath):
        #     net.load_state_dict(torch.load(self.parampath))  # 导入训练参数
        net.train()
        if self.iscuda == True:
            net = net.cuda()
        optimer = optim.Adam(net.parameters())
        con_lossfun = nn.BCELoss()  # 置信度损失函数
        off_lossfun = nn.MSELoss()  # 坐标的损失函数
        landoff_lossfun = nn.MSELoss()  # 坐标的损失函数
        for epoch in range(100000000):
            # 训练置信度
            if self.island:
                imgdata, con, offset, landoff = utils.GetBatch(
                    self.dataloader, True)
                landoff = Variable(landoff)
                if self.iscuda == True:
                    landoff = landoff.cuda()
            else:
                imgdata, con, offset = utils.GetBatch(self.dataloader, False)
            imgdata = Variable(imgdata)
            con = Variable(con)
            offset = Variable(offset)
            if self.iscuda == True:
                imgdata, con, offset = imgdata.cuda(), con.cuda(), offset.cuda(
                )

            if self.island:
                con_out, off_out, landoff_out = net(imgdata)
            else:
                con_out, off_out = net(imgdata)
            #计算置信度损失
            conn_mask = torch.lt(con, 2)
            conn_ = con[conn_mask]
            conn_out = con_out[conn_mask]
            con_loss = con_lossfun(conn_out, conn_)

            # 计算坐标偏移损失
            off_mask = torch.gt(con, 0)
            off_ = offset[off_mask[:, 0]]
            offf_out = off_out[off_mask[:, 0]]
            off_loss = off_lossfun(offf_out, off_)

            if self.island:
                #计算5个特征点的损失
                landoff_ = landoff[off_mask[:, 0]]
                landoff_out_ = landoff_out[off_mask[:, 0]]
                land_loss = landoff_lossfun(landoff_out_, landoff_)

                loss = con_loss + off_loss * 0.5 + land_loss
            else:
                loss = con_loss + off_loss * 0.5
            optimer.zero_grad()
            loss.backward()
            optimer.step()

            # 画图
            # plt.scatter(epoch, conloss, c='r', marker='.')
            # plt.scatter(epoch, offloss, c='y', marker='.')
            # plt.show()
            # plt.pause(1)
            if epoch % 10 == 0:
                conloss = con_loss.cpu().data.numpy()
                offloss = off_loss.cpu().data.numpy()
                loss_ = loss.cpu().data.numpy()
                if self.island:
                    landloss = land_loss.cpu().data.numpy()
                    print("epoch:", epoch, "loss:", loss_, "conloss:", conloss,
                          "offloss:", offloss, "landloss:", landloss)
                else:
                    print("epoch:", epoch, "loss:", loss_, "conloss:", conloss,
                          "offloss:", offloss)

                if True:
                    torch.onnx.export(
                        net,
                        imgdata,
                        r"C:\Users\liev\Desktop\myproject\face_recognition\params\yolonet.proto",
                        verbose=True)
                torch.save(net.state_dict(), self.parampath)  # 保存训练参数
            break