Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add ability to limit ITQ training memory usage #373

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Add ability to limit ITQ training memory usage #373

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
528dddc
Teach train_itq an option "max_descriptors"
predicative Oct 17, 2018
22ec85a
Add a release notes entry for train_itq's new option max_descriptors
predicative Oct 23, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions docs/release_notes/pending_release.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,11 @@ SMQTK Pending Release Notes
Updates / New Features
----------------------

Scripts
- `train_itq`
- Added an optional configuration property
``max_descriptors``. The descriptors used to train the ITQ
model are a random sample of the available descriptors.

Fixes
-----
Expand Down
30 changes: 27 additions & 3 deletions python/smqtk/bin/train_itq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,19 @@
be used to specify a sub-set of descriptors in the configured index to
train on. This only works if the stored descriptors' UUID is a type of
string.

The ``max_descriptors'' configuration property is optional and can be
used to cap the number of descriptors used to train the model. If
more descriptors are available than requested, they are randomly
subsampled.
"""

import logging
import os.path

import numpy
from six.moves import zip

from smqtk.algorithms.nn_index.lsh.functors.itq import ItqFunctor
from smqtk.representation import (
get_descriptor_index_impls,
Expand All @@ -32,6 +40,7 @@ def default_config():
"itq_config": ItqFunctor.get_default_config(),
"uuids_list_filepath": None,
"descriptor_index": plugin.make_config(get_descriptor_index_impls()),
"max_descriptors": None,
}


Expand All @@ -45,6 +54,7 @@ def main():
log = logging.getLogger(__name__)

uuids_list_filepath = config['uuids_list_filepath']
max_descriptors = config['max_descriptors']

log.info("Initializing ITQ functor")
#: :type: smqtk.algorithms.nn_index.lsh.functors.itq.ItqFunctor
Expand All @@ -63,12 +73,26 @@ def uuids_iter():
with open(uuids_list_filepath) as f:
for l in f:
yield l.strip()
uuids = uuids_iter()
log.info("Loading UUIDs list from file: %s", uuids_list_filepath)
d_iter = descriptor_index.get_many_descriptors(uuids_iter())
if max_descriptors:
uuids = list(uuids)
if max_descriptors < len(uuids):
log.info("Subsampling UUIDs (old count=%d, new count=%d)",
len(uuids), max_descriptors)
uuids = numpy.random.choice(uuids, max_descriptors, replace=False)
d_iter = descriptor_index.get_many_descriptors(uuids)
else:
d_length = len(descriptor_index)
log.info("Using UUIDs from loaded DescriptorIndex (count=%d)",
len(descriptor_index))
d_iter = descriptor_index
d_length)
if max_descriptors and max_descriptors < d_length:
log.info("Subsampling loaded DescriptorIndex (new count=%d)",
max_descriptors)
selected = numpy.random.permutation(numpy.arange(d_length) < max_descriptors)
d_iter = (d for d, s in zip(descriptor_index, selected) if s)
else:
d_iter = descriptor_index

log.info("Fitting ITQ model")
functor.fit(d_iter)
Expand Down