File size: 6,957 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
from __future__ import annotations
import logging
from copy import deepcopy
from dataclasses import dataclass, replace
from typing import Any, Mapping, Protocol, Sequence, TypeVar, runtime_checkable
import torch
from torch_scatter import scatter
T = TypeVar("T")
logger = logging.getLogger(__name__)
@runtime_checkable
class BatchedData(Protocol):
def replace(self: T, **vals: torch.Tensor) -> T:
"""Return a copy of self with some fields replaced with new values."""
def get_batch_idx(self, field_name: str) -> torch.LongTensor | None:
"""Get the batch index (i.e., which row belongs to which sample) for a given field.
For 'dense' type data, where every sample has the same shape and the first dimension is the
batch dimension, this method should return None. Mathematically,
returning None will be treated the same as returning a tensor [0, 1, 2, ..., batch_size - 1]
but I expect memory access in other functions to be more efficient if you return None.
"""
def get_batch_size(self) -> int:
"""Get the batch size."""
def device(self) -> torch.device:
"""Get the device of the batch."""
def __getitem__(self, field_name: str) -> torch.Tensor:
"""Get a field from the batch."""
def to(self: T, device: torch.device) -> T:
"""Move the batch to a given device."""
def clone(self: T) -> T:
"""Return a copy with all the tensors cloned."""
@dataclass
class SimpleBatchedData(BatchedData):
"""Implements BatchedData as a pair of mappings from field names to tensors."""
data: Mapping[str, Any]
batch_idx: Mapping[str, torch.LongTensor]
def replace(self, **vals: torch.Tensor) -> SimpleBatchedData:
"""Return a copy of self with some fields of self.data replaced with new values."""
return replace(self, data=self._updated_data(**vals))
def get_batch_idx(self, field_name: str) -> torch.LongTensor | None:
return self.batch_idx[field_name]
def _updated_data(self, **vals):
return dict(self.data, **vals)
def __getitem__(self, key):
return self.data[key]
def __contains__(self, key):
return key in self.data
def get_batch_size(self) -> int:
L = []
for k, v in self.batch_idx.items():
if v is None:
d = self.data[k]
L.append(d.shape[0] if isinstance(d, torch.Tensor) else len(d))
else:
if len(v) == 0:
logger.warning(f"Empty batch index for field {k}")
L.append(0)
else:
L.append(int(torch.max(v).item()) + 1)
return max(L)
@property
def device(self) -> torch.device:
# Beware, there are no checks that all values are on the same device
return next(v.device for v in self.data.values())
def to(self, device) -> "SimpleBatchedData":
"""Modify self in-place to move all tensors to the given device, and return self"""
if isinstance(self.data, dict):
for k in self.data.keys():
if isinstance(self.data[k], torch.Tensor):
self.data[k] = self.data[k].to(device)
if isinstance(self.batch_idx, dict):
for key in self.batch_idx.keys():
if self.batch_idx[key] is None:
continue
self.batch_idx[key] = self.batch_idx[key].to(device)
return self
def clone(self) -> SimpleBatchedData:
return SimpleBatchedData(
data={
k: v.clone() if isinstance(v, torch.Tensor) else deepcopy(v)
for k, v in self.data.items()
},
batch_idx={k: v.clone() if v is not None else None for k, v in self.batch_idx.items()},
)
def to_data_list(self) -> list[dict[str, torch.Tensor]]:
"""Converts this instance to a list of dictionaries, each of which corresponds to a single datapoint in
`batched_data`. The keys of the dictionaries match the keys of `batched_data`.
"""
batch_size = self.get_batch_size()
if batch_size == 0:
return []
def _unpack(k, i):
if self.batch_idx[k] is not None:
return self.data[k][self.batch_idx[k] == i]
elif isinstance(self.data[k], torch.Tensor):
return self.data[k][i : i + 1]
else:
return self.data[k][i]
return [{k: _unpack(k, i) for k in self.data.keys()} for i in range(batch_size)]
def collate_fn(
states: list[dict[str, Any]], dense_field_names: Sequence[str] = ()
) -> SimpleBatchedData:
"""
Combine a list of samples into a SimpleBatchedData object.
The association between the index in `states[i][k]` and a row in the `batched_data[k]` is
stored in `batched_data.batch_idx[k]`. If the `k` appears in
`dense_field_names`, `batched_data.batch_idx[k]` is `None` and the data is
simply stacked along the first dimension.
Non-tensor values are put into lists.
"""
assert states, "Cannot collate empty list"
concatenated_data = {}
batch_idx: dict[str, torch.Tensor | None] = {}
for k, v in states[0].items():
if isinstance(v, torch.Tensor):
concatenated_data[k] = torch.cat([x[k] for x in states], dim=0)
if k in dense_field_names:
if any(x[k].shape[0] != 1 for x in states):
raise ValueError(
f"First dimension should be batch dimension. Instead key {k} has shape {states[0][k].shape}"
)
batch_idx[k] = None
else:
batch_idx[k] = _construct_batch_idx(states, k)
else:
concatenated_data[k] = [x[k] for x in states]
batch_idx[k] = None
batch = SimpleBatchedData(data=concatenated_data, batch_idx=batch_idx)
if "edge_index" in batch.data:
batch = batch.replace(
edge_index=_batch_edge_index(
batch["edge_index"],
batch.batch_idx["atomic_numbers"],
batch.batch_idx["edge_index"],
),
)
return batch
def _batch_edge_index(edge_index, atom_batch_idx, edge_batch_idx):
num_atoms = scatter(torch.ones_like(atom_batch_idx), atom_batch_idx)
num_atoms_acc = torch.nn.functional.pad(torch.cumsum(num_atoms, 0)[:-1], [1, 0], "constant", 0)
return edge_index + num_atoms_acc[edge_batch_idx].unsqueeze(1)
def _construct_batch_idx(data_list: list[Any], field_name: str) -> torch.LongTensor:
"""Construct batch index tensor for one field."""
batch_size = len(data_list)
return torch.repeat_interleave(
torch.arange(0, batch_size),
torch.tensor([x[field_name].shape[0] for x in data_list]),
)
|