这是我参与「第三届青训营 -后端场」笔记创作活动的的第3篇笔记
为了提高实现效率和部署方便,选择将图片直接转化为前文的512长度向量存储,而不是下载图片本身,那么就涉及到如何高效下载数据集中url对应的图片,并实现边下载边用GPU处理,这需要用到多线程,以及共享变量加锁操作.
future 是一种对象，表示异步执行的操作。这个概念是 concurrent.futures模块和asyncio包的基础。
concurrent.futures 模块是Python3.2 引入的，对于Python2x 版本，Python2.5 以上的版本可以安装 futures 包来使用这个模块。
从Python3.4起，标准库中有两个为Future的类：concurrent.futures.Future 和 asyncio.Future。这两个类作用相同：两个Future类的实例都表示可能已经完成或未完成的延迟计算。
刚开始futures.ThreadPoolExecutor()参数的设置,因为没有完全了解多线程的机制,所以设置的值比较保守,以及一开始在jupyter中跑多线程程序,可能其对多线程机制支持不太好,一是速度慢,二是对大线程不太支持,容易卡死,后面切换为直接用python程序跑,提高参数,问题得到了解决.
一开始加锁的粒度太大,因为对加锁机制不太清楚,导致数据集中图片在下载到200000左右以后就会出现卡死现象,后边将加锁粒度减小,这个问题得到了有效缓解.
from torchvision import transforms
from torchvision import models
import time
import numpy as np
import faiss
from PIL import Image
import matplotlib.pyplot as plt
import requests
import io
from PIL import ImageFile
lock = threading.Lock()
def download_image(url):
global total
resp = requests.get(url, headers=fake_headers)
if resp.status_code == 200:
img = Image.open(io.BytesIO(resp.content))
if img.mode == 'RGB':
input_x = transform(img).unsqueeze(0)
feature = feature_extract(model, input_x.to(device))
feature = feature.data.cpu().detach().numpy()
lock.acquire()
feature_list = [str(f) for f in feature[0]]
with open('wukong/wukong_release/' + str(total) + '.txt', 'w', encoding='utf-8') as ff:
ff.write(url + '\n')
ff.write(" ".join(feature_list))
total += 1
lock.realease()
