def __init__(self, pnet_param, rnet_param, onet_param, 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()
# 定义transform为ToTensor
self.__image_transform = transforms.Compose([transforms.ToTensor()])
def __init__(self, pnet_param, rnet_param, onet_param, isCuda=False):
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()
# 装载网络训练结果
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()
# 将图片数据转换成NCHW
self.__image_transform = transforms.Compose(transforms.ToTensor())
#!/usr/bin/env python3.6
# -*- coding:utf-8 -*-
""" __author__ = "YYF"
__MTime__ = 18-11-26 上午11:46
"""
import Nets
import Train
if __name__ == '__main__':
net = Nets.ONet()
Trainer = Train.Trainer(net, './param/onet.pt',
r'/home/lievi/celeba_gen/48')
Trainer.train()