Feeding a model

Feeding images and their metadata to a model, a batch at a time.

For instance:

import vitrio as vio

particles = starfile.read("particles.star", always_dict=True)["particles"]
dataset = vio.data.Dataset({
        "image": vio.data.ImageSource(
                vio.data.parse_locations(particles["rlnImageName"]),
                data_type="float32"
        ),
        "defocus": particles[["rlnDefocusU", "rlnDefocusV"]].to_numpy(),
})
for batch in vio.data.DataLoader(dataset, 64, shuffle=True):
        ...

A Dataset is rows of named fields. A field is a source, which reads its values from files (ImageSource, PatchSource) or holds them in memory (ArraySource), or a column of metadata. A DataLoader goes through a dataset in batches, reading the next ones on a thread pool meanwhile, and hands them out as numpy arrays, tensors of PyTorch or arrays of JAX. The metadata is parsed by whoever reads it: vitrio takes what it was parsed into, and as_columns, parse_locations and make_locations turn that into fields.

Needs numpy. PyTorch and JAX are imported only to hand batches to them. This file is the whole public surface of the subpackage.

class vitrio.data.Dataset(fields, index='index')

Rows of named fields: images, patches, and the metadata that goes with them.

A field is a source, such as an ImageSource, or a column, which is anything numpy takes as an array of a value for each row: an array of CTF parameters, a column of a DataFrame. Every field has a value for each row.

Indexing a dataset with an integer reads one row, and with a slice or an array of integers reads a batch of rows into an array for each field. Each batch also holds the index of each of its rows, under index. So it works as it is with the DataLoader of PyTorch, which reads a batch at a time when its sampler draws batches:

torch.utils.data.DataLoader(
        dataset,
        sampler=torch.utils.data.BatchSampler(
                torch.utils.data.RandomSampler(dataset), 64, False
        ),
        batch_size=None
)

A dataset only describes its rows, and holds nothing that reads them. Each read is given the loader it goes through, a thread pool with its open readers and maybe a scratch, and a read given none goes through a loader shared by the process. Indexing always does. DataLoader of vitrio.data takes a loader of its own, and needs no worker process.

Parameters:

• fields – The fields by name, each a source or a column.

• index – The name the indices of the rows of a batch go under. None leaves them out.

Raises:

ValueError – If there are no fields, they are not all as long, or one has the name index is given.

property fields

The fields by name.

property index

The name the indices of the rows go under, if anywhere.

read_async(rows, allocator=None, loader=None)

Start reading a batch of rows.

Parameters:

• rows – The rows, as integers, negative ones counting from the end, or as a mask of a boolean for each row. A row may be asked for more than once.

• allocator – What makes the array of each field, given its shape and its numpy type. Defaults to numpy.empty. An allocator of pinned memory lets PyTorch copy the batch to a GPU without staging it.

• loader – Where reads of files are dispatched. Defaults to the loader shared by the process.

Returns:

The batch, once it is complete.

Return type:

PendingBatch

subset(rows)

Get a dataset of only some of the rows, in the order asked for.

It shares the sources of this one. Its indices count its own rows.

Parameters:

rows – The rows, named as read_async takes them.

preload(names=None, loader=None)

Get a dataset whose sources have all been read into memory.

Going through it afterwards reads no file.

Parameters:

• names – The fields to read. Defaults to every source not already in memory.

• loader – Where reads of files are dispatched. Defaults to the loader shared by the process.

class vitrio.data.DataLoader(dataset, batch_size=1, shuffle=False, *, batch_sampler=None, loader=None, drop_last=False, pad_last=False, seed=None, prefetch=4, output='numpy', device=None, sharding=None, pin_memory=None, allocator=None, transform=None, valid='valid')

Goes through a dataset in batches, reading the next ones meanwhile.

Its arguments are those of PyTorch’s DataLoader where they mean the same, but it starts no worker process: the images are read on the thread pool of its loader, prefetch batches at once, straight into the arrays they are handed out in. Each iteration is an epoch.

The loader is what reads: its thread pool, its open readers and, when it was made with one, a scratch that serves the epochs after the first. Loaders that are given one loader share its thread pool.

A batch is a dict with an array for each field of the dataset, the first axis of each running along the rows of the batch. It is handed out as numpy arrays, as tensors of PyTorch or as arrays of JAX. A field that does not hold numbers stays a numpy array.

JAX compiles a function again for every shape it is given, so a last batch with fewer rows costs a compilation. drop_last leaves it out, and pad_last fills it up with rows of its own and says which rows are real in a field of booleans.

Every batch is read into arrays of its own, which PyTorch and JAX may use without a copy, so a batch that has been handed out is never written again.

Parameters:

• dataset – What to go through.

• batch_size – How many rows a batch has.

• shuffle – Whether to draw the rows in an order of each epoch’s own.

• batch_sampler – What draws the rows of each batch, instead of a BatchSampler made of batch_size, shuffle, drop_last and seed. Any iterable of sequences of rows, such as a BatchSampler of PyTorch over its DistributedSampler.

• loader – Where reads of files are dispatched, such as vio.loader(workers=8, scratch=scratch). Defaults to the loader shared by the process.

• drop_last – Whether to leave out the last batch when it has fewer rows.

• pad_last – Whether to fill the last batch up to batch_size rows by repeating its own, and give every batch a field of booleans, valid, false for the rows that repeat.

• seed – What decides the orders of the epochs. Defaults to a seed of the loader’s own.

• prefetch – How many batches to have under way at once.

• output – What to hand the batches out as: ‘numpy’, ‘torch’, ‘jax’, or a callable given each batch as numpy arrays.

• device – Where to copy the batches to: a device of PyTorch for ‘torch’, or of JAX for ‘jax’. Defaults to leaving them in host memory for PyTorch, and to the default device of JAX.

• sharding – How JAX lays each field out across its devices, for ‘jax’ only. A sharding over the devices of several processes makes each batch the share of this process.

• pin_memory – Whether to read into pinned memory, for ‘torch’ only. Defaults to doing so when device is a GPU of CUDA.

• allocator – What makes the arrays of a batch, as Dataset.read_async takes it. Defaults to numpy.empty.

• transform – What to do with each batch once it is handed out as output says, such as normalizing it on the GPU.

• valid – The name of the field pad_last adds.

Raises:

ValueError – If a batch sampler is given with any of what it replaces, drop_last with pad_last, output is none of the above, or what is given does not apply to it.

class vitrio.data.BatchSampler(size, batch_size, shuffle=False, drop_last=False, seed=None, shard=(0, 1))

Draws the rows of a dataset in batches, once each per epoch.

Each iteration is an epoch. When it shuffles, every epoch has an order of its own, which the seed and the number of the epoch decide, so two samplers given one seed draw the same rows.

A sampler can draw only one shard of the rows, for one of several processes that each read a part of every batch, as with PyTorch’s DistributedSampler: the rows of an epoch are dealt out in turn, after the first ones are repeated until every shard has as many. The shards must be given one seed, so that they deal out the same order.

Parameters:

• size – How many rows there are.

• batch_size – How many rows a batch has, the last one maybe fewer. With shards, the rows of a batch of one shard.

• shuffle – Whether to draw the rows in an order of each epoch’s own. Otherwise they are drawn in their order.

• drop_last – Whether to leave out the last batch when it has fewer rows.

• seed – What decides the orders. Defaults to a seed of the sampler’s own.

• shard – Which shard to draw, and of how many.

Raises:

ValueError – If the batch size is not positive, or the shard is not one of its count.

property epoch

The epoch the next iteration draws.

set_epoch(epoch)

Choose the epoch the next iteration draws.

order(epoch)

The rows of the shard in the order an epoch draws them.

Fields

class vitrio.data.Source(*args, **kwargs)

The values of one field of a dataset, a value for each row.

shape and dtype are those of the value of one row, so the values of n rows are read into an array of shape (n, *shape).

property shape

The shape of the value of one row.

property dtype

The data type of the values.

read_async(loader, rows, out)

Start reading the values of some rows.

Parameters:

• loader – Where reads of files are dispatched.

• rows – The rows to read, as integers within the length of the source.

• out – Where the value of each row lands, in the order of rows. Its shape is (len(rows), *shape) and its type dtype.

Returns:

What to wait for before out holds the values. It may be empty, when the values are in out already.

subset(rows)

Get the source of only these rows, in this order.

class vitrio.data.ImageSource(locations, data_type=None, readers=None, check=True)

Images or volumes read from files, a location for each row.

A row may have several locations instead, as many for every row: the tilts of a particle, the halves of a movie, an image and its target. Its value is then those images stacked, and shape has the number of locations of a row in front of the shape of one image.

Every location has to name an image of the same shape. The files are asked for their shape when the source is made, each file once, and a location that does not fit is reported then.

Parameters:

• locations – The location of the image of each row, or a row of locations for each row: ImageLocation objects, or keys of whole files.

• data_type – The type to read the values as. Defaults to the one the first file holds.

• readers – Where files are asked for their shape. Defaults to a provider that opens the file at the path a key is.

• check – Whether to ask every file for its shape, and every stack whether it holds the indices asked of it. Otherwise only the first file is asked.

Raises:

ValueError – If two locations name images of different shapes, an index is past the end of its stack, or there are no locations.

property locations

The location of each row, as an array of ImageLocation objects.

read_async(loader, rows, out)

Start reading the values of some rows, as Source.read_async.

subset(rows)

Get the source of only these rows, in this order.

class vitrio.data.PatchSource(images, centres, shape, data_type=None, fill=0, readers=None, check=True)

Patches cropped out of larger images, a patch for each row.

Each row names the image its patch is cropped from and the centre of the patch, which lands at index extent // 2 within it, as RELION has it. A patch reaching past the edge of its image is read as far as the image goes, and the rest of it is fill. The rows of a batch that share an image are read in one call.

Particles picked on micrographs are patches, and so are the boxes of a tomogram around its picked objects, or around random centres.

Parameters:

• images – The image each row is cropped from: an ImageLocation or the key of a file that holds one image or volume.

• centres – The centre of each patch, a row of integers for each row, slowest axis first, as (y, x) or (z, y, x).

• shape – The shape of one patch.

• data_type – The type to read the values as. Defaults to the one the first image holds.

• fill – The value of what a patch has beyond its image.

• readers – Where files are asked for their shape. Defaults to a provider that opens the file at the path a key is.

• check – Whether to ask every image for its shape. Otherwise only the first is asked.

Raises:

ValueError – If the images, centres and patches are not of one rank, there are not as many centres as images, a centre is negative, or there are no rows.

property images

The images the patches are cropped from, each once.

property centres

The centre of the patch of each row.

read_async(loader, rows, out)

Start reading the values of some rows, as Source.read_async.

subset(rows)

Get the source of only these rows, in this order.

class vitrio.data.ArraySource(array)

Values already in memory, the first axis of an array running along the rows.

Parameters:

array – The values. Anything numpy takes as an array, which it uses without a copy where it can.

property array

The values.

read_async(loader, rows, out)

Start reading the values of some rows, as Source.read_async.

subset(rows)

Get the source of only these rows, in this order.

class vitrio.data.PendingBatch(values, pending)

A batch under way. Its arrays are handed out once it is complete.

property is_ready

Whether every read of the batch is done.

get()

Wait for the batch.

Returns:

The value of each field for each row of the batch, in the order its rows were asked for.

Return type:

dict[str, np.ndarray]

Raises:

• Whatever a read of the batch raised. Every read is waited for –

• first. –

Tables

vitrio.data.as_columns(table)

Split a table into one array for each of its columns.

Parameters:

table – A pandas DataFrame, a mapping of names to array-likes or a numpy structured array.

Returns:

The columns by name, in the order of the table. Each holds a row of the table at each index of its first axis.

Return type:

dict[str, np.ndarray]

Raises:

• TypeError – If table is none of the above.

• ValueError – If the columns are not all as long.

vitrio.data.make_locations(keys, indices=None, root=None)

Make locations out of keys and the indices in their stacks.

That is how cryoSPARC states where a particle is: its blob/path and blob/idx fields, the index counting from zero. The > that cryoSPARC writes in front of a path relative to its project is left out.

Parameters:

• keys – The file of each location, as strings, bytes or paths.

• indices – The index of each location in its stack, counting from zero. Defaults to locations of whole files.

• root – The directory a relative key is relative to. Defaults to leaving it relative, which is to the working directory.

Returns:

The locations, in the order of keys.

Return type:

list[ImageLocation]

Raises:

ValueError – If keys and indices are not as long.

vitrio.data.parse_locations(texts, root=None)

Parse locations written as text, as RELION writes them.

A location is written “index@key”, the index counting from one, or “key” for a whole file. That is how the rlnImageName and rlnMicrographName columns of a STAR file have them.

Parameters:

• texts – The locations, as strings or bytes.

• root – The directory a relative key is relative to. Defaults to leaving it relative, which is to the working directory.

Returns:

The locations, in the order of texts.

Return type:

list[ImageLocation]

Raises:

ValueError – If a text is not a location.