def __init__(self, model_path):
self.model = bodypose_model()
if torch.cuda.is_available():
self.model = self.model.cuda()
model_dict = util.transfer(self.model, torch.load(model_path))
self.model.load_state_dict(model_dict)
self.model.eval()
def __init__(self, model_path):
self.model = bodypose_model() # 创建模型
if torch.cuda.is_available():
self.model = self.model.cuda()
model_dict = util.transfer(self.model,
torch.load(model_path)) # 参数格式转换
self.model.load_state_dict(model_dict) # 导入参数
self.model.eval() # 使网络中bn层和dropout层失效。
def __init__(self, model_path):
self.model = handpose_model()
self.cuda_true = False
if torch.cuda.is_available():
self.model = self.model.cuda()
self.cuda_true = True
model_dict = util.transfer(self.model, torch.load(model_path))
self.model.load_state_dict(model_dict)
self.model.eval()
out4_1 = self.model4_1(out4)
out4_2 = self.model4_2(out4)
out5 = torch.cat([out4_1, out4_2, out1], 1)
# Stage5
out5_1 = self.model5_1(out5)
out5_2 = self.model5_2(out5)
out6 = torch.cat([out5_1, out5_2, out1], 1)
# Stage6
out6_1 = self.model6_1(out6)
out6_2 = self.model6_2(out6)
return out6_1, out6_2
if __name__ == "__main__":
import torch
from torchsummary import summary
from util import transfer
PRETRAIN_URL = "https://www.dropbox.com/sh/7xbup2qsn7vvjxo/AABaYNMvvNVFRWqyDXl7KQUxa/body_pose_model.pth?dl=1"
state_dict = torch.hub.load_state_dict_from_url(PRETRAIN_URL)
model = bodypose_model()
state_dict = transfer(model, state_dict)
model.load_state_dict(state_dict)
model.eval()
summary(model, (3, 512, 512), device="cpu")
