def __init__(self, pnet_param="./param_0/pnet.pt", rnet_param="./param_0/rnet.pt",onet_param="./param_0/onet.pt", isCuda=True): self.isCuda = isCuda self.pnet = nets.PNet() # 创建实例变量,实例化P网络 self.rnet = nets.RNet() self.onet = nets.ONet() if self.isCuda: self.pnet.cuda() # 给P网络加速 self.rnet.cuda() self.onet.cuda() self.pnet.load_state_dict(torch.load(pnet_param)) # 把训练好的权重加载到P网络中 self.rnet.load_state_dict(torch.load(rnet_param)) self.onet.load_state_dict(torch.load(onet_param)) # self.pnet.eval() # 训练网络里有BN(批归一化时),要调用eval方法,使用是不用BN,dropout方法 # self.rnet.eval() # self.onet.eval() self.__image_transform = transforms.Compose([ transforms.ToTensor() ]) # 图片数据类型转换
def __init__(self, pnet_param="./param/pnet.pt", rnet_param="./param/rnet.pt", onet_param="./param/onet.pt", isCuda=True): self.isCuda = isCuda self.pnet = nets.PNet() self.rnet = nets.RNet() self.onet = nets.ONet() if self.isCuda: self.pnet.cuda() self.rnet.cuda() self.onet.cuda() self.pnet.load_state_dict(torch.load(pnet_param)) self.rnet.load_state_dict(torch.load(rnet_param)) self.onet.load_state_dict(torch.load(onet_param)) self.pnet.eval() self.rnet.eval() self.onet.eval() self.__image_transform = transforms.Compose([ transforms.ToTensor() ])
def __init__(self, pnet_param=r"./param_0/pnet.pt", rnet_param=r"./param_0/rnet.pt", onet_param=r"./param_0/onet.pt", isCuda=True): self.isCuda = isCuda # 创建实例变量,实例化网络 self.pnet = nets.PNet() self.rnet = nets.RNet() self.onet = nets.ONet() # 加入cuda加速训练 if self.isCuda: self.pnet.cuda() self.rnet.cuda() self.onet.cuda() # 把训练好的权重加载到p网络中 self.pnet.load_state_dict(torch.load(pnet_param)) self.rnet.load_state_dict(torch.load(rnet_param)) self.onet.load_state_dict(torch.load(onet_param)) # 调用测试模式,因为网络中含有BN self.pnet.eval() self.rnet.eval() self.onet.eval() # 图片数据类型转换 self.__image_transform = transforms.Compose([transforms.ToTensor()])
def __init__(self, pnet_param="./param6/p_net.pth", rnet_param="./param6/r_net.pth", onet_param="./param6/o_net.pth", isCuda=True): self.isCuda = isCuda self.pnet = nets.PNet() self.rnet = nets.RNet() self.onet = nets.ONet() if self.isCuda: self.pnet.cuda() self.rnet.cuda() self.onet.cuda() self.pnet.load_state_dict(torch.load(pnet_param, map_location="cuda")) self.rnet.load_state_dict(torch.load(rnet_param, map_location="cuda")) self.onet.load_state_dict(torch.load(onet_param, map_location="cuda")) self.pnet.eval() # 批归一化, 使用之前训练的Batchnormal self.rnet.eval() self.onet.eval() self._image_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ])
def __init__(self, pnet_param="./param/p_net.pth", rnet_param="./param/r_net.pth", onet_param="./param/o_net.pth", isCuda=False): self.isCuda = isCuda self.pnet = nets.PNet() self.rnet = nets.RNet() self.onet = nets.ONet() if self.isCuda: self.pnet.cuda() self.rnet.cuda() self.onet.cuda() self.pnet.load_state_dict(torch.load(pnet_param)) self.rnet.load_state_dict(torch.load(rnet_param)) # self.onet.load_state_dict(torch.load(onet_param, map_location='cpu')) self.pnet.eval() self.rnet.eval() self.onet.eval() self.__image_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])
def __init__(self, pnet_param="./param/pnet.pt", rnet_param="./param/rnet.pt", onet_param="./param/onet.pt", isCuda=False): self.isCuda = isCuda self.pnet = nets.PNet() self.rnet = nets.RNet() self.onet = nets.ONet() if self.isCuda: self.pnet.cuda() self.rnet.cuda() self.onet.cuda() self.pnet.load_state_dict(torch.load(pnet_param, map_location= "cpu" )) self.rnet.load_state_dict(torch.load(rnet_param, map_location= "cpu" )) self.onet.load_state_dict(torch.load(onet_param, map_location= "cpu" )) self.__image_transform = transforms.Compose([transforms.ToTensor()]) #NCHW and normalized
def __init__(self, pnet_path, rnet_path, onet_path, isCuda=True): self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.isCuda = isCuda self.pnet = nets.PNet().to(self.device) self.rnet = nets.RNet().to(self.device) self.onet = nets.ONet().to(self.device) self.pnet.load_state_dict(torch.load(pnet_path)) self.rnet.load_state_dict(torch.load(rnet_path)) self.onet.load_state_dict(torch.load(onet_path)) self.pnet.eval() self.rnet.eval() self.onet.eval() self.trans = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ])
def __init__(self, p_net_param=r"./save_params/p_net.pth", r_net_param=r"./save_params/r_net.pth", o_net_param="./save_params/o_net.pth", isCuda=False): self.isCuda = isCuda self.p_net = nets.PNet().to(DEVICE) self.r_net = nets.RNet().to(DEVICE) self.o_net = nets.ONet().to(DEVICE) self.p_net.load_state_dict(torch.load(p_net_param)) self.r_net.load_state_dict(torch.load(r_net_param)) self.o_net.load_state_dict(torch.load(o_net_param)) self.p_net.eval() self.r_net.eval() self.o_net.eval() self.__image_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 对于图像这个是有经验值的 ])
def __init__(self, pnet_path, rnet_path, onet_path, softnms=False, thresholds=None, factor=0.709): if thresholds is None: thresholds = [0.6, 0.6, 0.95] self.thresholds = thresholds self.factor = factor self.softnms = softnms self.pnet = nets.PNet().to(device) self.rnet = nets.RNet().to(device) self.onet = nets.ONet().to(device) self.pnet.load_state_dict(torch.load(pnet_path)) self.rnet.load_state_dict(torch.load(rnet_path)) self.onet.load_state_dict(torch.load(onet_path)) self.pnet.eval() self.rnet.eval() self.onet.eval() self.img_transfrom = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ])
# 训练P网络 import nets import train if __name__ == '__main__': net = nets.PNet() trainer = train.Trainer(net, './param/pnet.pt', r"C:\celeba_2\12") # 网络、保存参数、训练数据 trainer.train() # 调用训练方法