def __init__(self,
pnet_param="E:/save_path/12/param/pnet_epoch_99.pt",
rnet_param="E:/save_path/24/param/pnet_epoch_99.pt",
onet_param="E:/save_path/48/param/pnet_epoch_99.pt",
isCuda=True):
self.isCuda = isCuda
self.pnet = net.PNet()
self.rnet = net.RNet()
self.onet = net.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="./train_01/pnet.pt",
rnet_param="./train_01/rnet.pt",
onet_param="./train_01/onet.pt",
isCuda=True):
self.isCuda = isCuda
'创建实例变量，实例化网络'
self.pnet = net.PNet()
self.rnet = net.RNet()
self.onet = net.ONet()
if self.isCuda:
'给网络放到CUDA上运行'
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))
'把训练网络里有BN(批归一化时)，要调用eval方法，使用是不用BN, dropout方法'
self.pnet.eval()
self.rnet.eval()
self.onet.eval()
'把图片数据类型转换,归一化及换轴操作'
self.__image_transform = transforms.Compose([transforms.ToTensor()])
def __init__(self):
"Device"
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
"Net"
self.pnet = net.PNet().to(self.device)
self.rnet = net.RNet().to(self.device)
self.onet = net.ONet().to(self.device)
if not os.path.exists("./params"):
raise FileNotFoundError
self.pnet.load_state_dict(torch.load("./params/pnet.pkl"))
self.rnet.load_state_dict(torch.load("./params/rnet.pkl"))
def __init__(self, para_p, para_r, para_o, test_img):
self.device = torch.device("cuda")
self.test_img = test_img
self.image = Image.open(test_img) # 用于抠图输入下一层
self.img = Image.open(test_img) # 复制图片用于图像金字塔
self.pnet = net.PNet().to(self.device)
self.pnet.load_state_dict(torch.load(para_p))
self.pnet.eval()
self.rnet = net.RNet().to(self.device)
self.rnet.load_state_dict(torch.load(para_r))
self.rnet.eval()
self.onet = net.ONet().to(self.device)
self.onet.load_state_dict(torch.load(para_o))
self.onet.eval()
def __init__(self, test_img):
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self.image = Image.fromarray(np.uint8(
test_img[:, :, [2, 1, 0]])) # for croped transform cv2 to PIL
self.img = self.image # for pyramid
self.pnet = net.PNet().to(self.device)
self.pnet.load_state_dict(torch.load("./params/pnet.pkl"))
self.pnet.eval()
self.rnet = net.RNet().to(self.device)
self.rnet.load_state_dict(torch.load("./params/rnet.pkl"))
self.rnet.eval()
self.onet = net.ONet().to(self.device)
self.onet.load_state_dict(torch.load("./params/onet.pkl"))
self.onet.eval()
def __init__(self, mode: str):
"Device"
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
"Net"
self.name = mode.lower() + "net"
if mode == "P" or mode == "p":
self.net = net.PNet().to(self.device)
self.train = dataset.P_TRAIN
self.test = dataset.P_TEST
elif mode == "R" or mode == "r":
self.net = net.RNet().to(self.device)
self.train = dataset.R_TRAIN
self.test = dataset.R_TEST
elif mode == "O" or mode == "o":
self.net = net.ONet().to(self.device)
self.train = dataset.O_TRAIN
self.test = dataset.O_TEST
if os.path.exists("./params/{}.pkl".format(self.name)):
print("MODE: {} >>> Loading ... ...".format(mode.upper()))
self.net.load_state_dict(
torch.load("./params/{}.pkl".format(self.name)))
else:
print("MODE: {} >>> Initing ... ...".format(mode.upper()))
self.net.apply(functions.weights_init)
"Optimize"
self.opt = optim.Adam(self.net.parameters())
"Loss"
self.Closs = nn.BCEWithLogitsLoss()
self.Hloss = nn.BCEWithLogitsLoss(reduction='none')
self.Oloss = nn.MSELoss()
self.LMloss = nn.MSELoss()
"Tensorboard"
self.draw = SummaryWriter(
log_dir="./runs/{}_runs".format(mode.lower()))
def __init__(self, mode: str, batch_size: int):
"Device Config"
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
"Model Config"
self.mode = mode
if mode == "P" or mode == "p":
self.size = 12
self.threshold = 0.9
self.net = net.PNet().to(self.device)
elif mode == "R" or mode == "r":
self.size = 24
self.threshold = 0.99
self.net = net.RNet().to(self.device)
elif mode == "O" or mode == "o":
self.size = 48
self.threshold = 0.999
self.net = net.ONet().to(self.device)
if len(os.listdir("./params")) > 3:
print("MODE: {} >>> Loading ... ...".format(mode))
self.net.load_state_dict(torch.load("./params/{}net.pkl".format(mode.lower())))
"Dataloader Config"
self.train = data.DataLoader(dataset.choice(mode.lower()), batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True)
self.test = data.DataLoader(dataset.choice("{}v".format(mode.lower())), batch_size=512, shuffle=True, drop_last=True)
"Optim Config"
optimize = optim.SGD(self.net.parameters(), lr=3e-4, momentum=0.9)
self.lookahead = Lookahead(optimize, k=5, alpha=0.5)
# self.lr = optim.lr_scheduler.CosineAnnealingLR(self.lookahead, T_max=1, eta_min=1e-5, last_epoch=-1)
"Loss Config"
self.loss_confi = nn.BCELoss()
self.loss_resolve = nn.BCELoss(reduction="none")
self.loss_offset = nn.MSELoss()
"Show Config"
self.summarywriter = SummaryWriter(log_dir="./runs/{}_runs".format(mode.lower()))
def __init__(
self,
p_net_param=r"./params_p/p_net_30.pth",
# r_net_param="./params_r/r_net_13_s.pth",
# 2(7),3(10),5(7),6(8),7(10),10(9),12(8), -13(5),15(8),21(7)
r_net_param=r"./params_r/r_net_53_s.pth",
o_net_param="./params_o/o_net_35.pth",
isCuda=False):
self.isCuda = isCuda
self.p_net = net.PNet().to(DEVICE)
self.r_net = net.RNet().to(DEVICE)
self.o_net = net.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]) # 对于图像这个是有经验值的
])
import net
import train
if __name__ == '__main__':
ne = net.ONet()
ne.train()
trainer = train.Trainer(ne, 'E:/save_path/48/param/', r"E:\save_path\48")
trainer.trainOnet()
''
'训练O网络'
import net
from train_01 import train
if __name__ == '__main__':
net = net.ONet()
trainer = train.Trainer(net, r'E:\MTCNN\June\cebela_01\48', 'onet.pt')
trainer.train()