首页 > 其他 > 详细

第一章基本操作-强大的hub模块

时间:2020-04-07 01:51:52      阅读:196      评论:0      收藏:0      [点我收藏+]

使用torch的hub模块载入模型,输入数据进行模型的结果输出,对输出的结果做可视化处理 

## GITHUB https://github.com/pytorch/hub 

import torch 

model = torch.hub.load(pytorch/vision:v0.4.2, deeplabv3_resnet101, pretrained=True)
model.eval() 

print(torch.hub.list(pytorch/vision:v0.4.2))

#数据载入,获得图片 
import urllib 
url, filename = ("https://github.com/pytorch/hub/raw/master/dog.jpg", "dog.jpg")
try:
    urllib.URLopener().retrieve(url, filename)
except:
    urllib.request.urlretrieve(url, filename)

from PIL import Image 
from torchvision import transforms 

input_image = Image.open(filename)
#构建处理图片的函数 
preprocess = transforms.Compose(
    [
        transforms.ToTensor(), 
        transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]), 
    ] 
)
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # 产生一个样本 

if torch.cuda.is_available():
    input_batch = input_batch.to("cuda")
    model.to("cuda")

with torch.no_grad():
    output = model(input_batch)[out][0]

output_predictions = output.argmax(0)

palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
colors = torch.as_tensor([i for i in range(21)])[:, None] * colors 
colors = (colors % 255).numpy().astype("uint8")

r = Image.fromarray(output_predictions.bytes().cpu().numpy()).resize(input_image.size)
r.putpalette(colors)

import matplotlib.pyplot as plt 
plt.show(r)

 

第一章基本操作-强大的hub模块

原文:https://www.cnblogs.com/my-love-is-python/p/12650573.html

(1)
(1)
   
举报
评论 一句话评论(0
关于我们 - 联系我们 - 留言反馈 - 联系我们:wmxa8@hotmail.com
© 2014 bubuko.com 版权所有
打开技术之扣,分享程序人生!