def __init__(self, frames):
self.model_I = Inet()
self.model_P = Pnet()
self.model_dir_inet = os.path.join(parent_folder15, 'models', 'I_e290.pth')
self.model_dir_pnet = os.path.join(parent_folder15, 'models', 'P_e290.pth')
if torch.cuda.is_available():
print('Using GPU')
self.model_I = nn.DataParallel(self.model_I) # 单服务器多GPU
self.model_P = nn.DataParallel(self.model_P)
self.model_I.cuda() # 将模型复制到gpu,默认是cuda('0'),即跳转到第一个gpu
self.model_P.cuda()
# load_state_dict加载static_dict-字典对象,将每一层与它的对应参数建立映射关系,只有参数可以训练的layer才会被保存
self.model_I.load_state_dict(torch.load(self.model_dir_inet))
self.model_P.load_state_dict(torch.load(self.model_dir_pnet))
else:
print('Using CPU')
self.model_I.load_state_dict(load_UnDP(self.model_dir_inet))
self.model_P.load_state_dict(load_UnDP(self.model_dir_pnet))
self.model_I.eval() # turn-off BN
self.model_P.eval() # turn-off BN
self.frames = frames.copy()
self.num_frames, self.height, self.width = self.frames.shape[:3]
print('num_frames:{} height:{} width:{}'.format(self.num_frames, self.height, self.width))
self.init_variables(self.frames)
def __init__(self, load_pretrain=True):
self.model_I = Inet()
self.model_P = Pnet()
if torch.cuda.is_available():
print('Using GPU')
self.model_I = nn.DataParallel(self.model_I)
self.model_P = nn.DataParallel(self.model_P)
self.model_I.cuda()
self.model_P.cuda()
if load_pretrain:
self.model_I.load_state_dict(torch.load('I_e290.pth'))
self.model_P.load_state_dict(torch.load('P_e290.pth'))
else:
print('Using CPU')
if load_pretrain:
self.model_P.load_state_dict(load_UnDP('P_e290.pth'))
self.model_I.load_state_dict(load_UnDP('I_e290.pth'))
self.eval()
def __init__(self, frames):
self.model_I = Inet()
self.model_P = Pnet()
if torch.cuda.is_available():
print('Using GPU')
self.model_I = nn.DataParallel(self.model_I)
self.model_P = nn.DataParallel(self.model_P)
self.model_I.cuda()
self.model_P.cuda()
self.model_I.load_state_dict(torch.load('I_e290.pth'))
self.model_P.load_state_dict(torch.load('P_e290.pth'))
else:
print('Using CPU')
self.model_I.load_state_dict(load_UnDP('I_e290.pth'))
self.model_P.load_state_dict(load_UnDP('P_e290.pth'))
self.model_I.eval() # turn-off BN
self.model_P.eval() # turn-off BN
self.frames = frames.copy()
self.num_frames, self.height, self.width = self.frames.shape[:3]
self.init_variables(self.frames)