Coverage for python/lsst/meas/transiNet/rbTransiNetInterface.py: 97%
60 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-01 09:00 +0000
1# This file is part of meas_transiNet.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22__all__ = ["RBTransiNetInterface", "CutoutInputs"]
24import dataclasses
25import math
27import numpy as np
28import torch
30import lsst.utils.logging
32from .modelPackages.nnModelPackage import NNModelPackage
35@dataclasses.dataclass(frozen=True, kw_only=True)
36class CutoutInputs:
37 """Science/template/difference cutouts of a single object plus other
38 metadata.
39 """
40 science: np.ndarray
41 template: np.ndarray
42 difference: np.ndarray
44 label: bool = None
45 """Known truth of whether this is a real or bogus object."""
48class RBTransiNetInterface:
49 """ The interface between the LSST AP pipeline and a trained pytorch-based
50 RBTransiNet neural network model.
52 Parameters
53 ----------
54 task : `lsst.meas.transiNet.RBTransiNetTask`
55 The task that is using this interface: the 'left side'.
56 model_package_name : `str`
57 Name of the model package to load.
58 package_storage_mode : {'local', 'neighbor'}
59 Storage mode of the model package
60 device : `str`
61 Device to load and run the neural network on, e.g. 'cpu' or 'cuda:0'
62 """
64 def __init__(self, task, device='cpu'):
65 self.task = task
67 # in case the model package name is not set at this stage, it is not
68 # needed (e.g. in butler mode).
69 self.model_package_name = task.config.modelPackageName or 'N/A'
71 self.package_storage_mode = task.config.modelPackageStorageMode
72 self.device = device
73 self.init_model()
75 def init_model(self):
76 """Create and initialize an NN model
77 """
79 if self.package_storage_mode == 'butler' and self.task.butler_loaded_package is None: 79 ↛ 80line 79 didn't jump to line 80 because the condition on line 79 was never true
80 raise RuntimeError("RBTransiNetInterface is trying to load a butler-mode NN model package, "
81 "but the RBTransiNetTask has not passed down a preloaded payload.")
83 model_package = NNModelPackage(model_package_name=self.model_package_name,
84 package_storage_mode=self.package_storage_mode,
85 butler_loaded_package=self.task.butler_loaded_package)
86 self.model = model_package.load(self.device)
87 self.version = str(model_package.metadata['version'])
89 # Put the model in evaluation mode instead of training model.
90 self.model.eval()
92 def input_to_batches(self, inputs, batchSize):
93 """Convert a list of inputs to a generator of batches.
95 Parameters
96 ----------
97 inputs : `list` [`CutoutInputs`]
98 Inputs to be scored.
100 Returns
101 -------
102 batches : `generator`
103 Generator of batches of inputs.
104 """
105 for i in range(0, len(inputs), batchSize):
106 yield inputs[i:i + batchSize]
108 def prepare_input(self, inputs):
109 """Convert inputs from numpy arrays, etc. to a torch.tensor blob.
111 Parameters
112 ----------
113 inputs : `list` [`CutoutInputs`]
114 Inputs to be scored.
116 Returns
117 -------
118 blob
119 Prepared torch tensor blob to run the model on.
120 labels
121 Truth labels, concatenated into a single list.
122 """
123 cutoutsList = []
124 labelsList = []
125 for inp in inputs:
126 # Convert each cutout to a torch tensor
127 template = torch.from_numpy(inp.template)
128 science = torch.from_numpy(inp.science)
129 difference = torch.from_numpy(inp.difference)
131 # Stack the components to create a single blob
132 # dimensions should be 3 x width x height
133 singleBlob = torch.stack([difference, science, template], dim=0)
135 # And append them to the temporary list
136 cutoutsList.append(singleBlob)
138 labelsList.append(inp.label)
140 blob = torch.stack(cutoutsList)
142 return blob, labelsList
144 def infer(self, inputs):
145 """Return the score of this cutout.
147 Parameters
148 ----------
149 inputs : `list` [`CutoutInputs`]
150 Inputs to be scored.
152 Returns
153 -------
154 scores : `numpy.array`
155 Float scores for each element of ``inputs``.
156 """
158 # Handle empty inputs gracefully.
159 if not inputs:
160 return np.array([])
162 # Convert the inputs to batches.
163 # TODO: The batch size is set to 64 for now. Later when
164 # deploying parallel instances of the task, memory limits
165 # should be taken into account, if necessary.
166 batch_size = 64
167 batches = self.input_to_batches(inputs, batchSize=batch_size)
169 # Log every 10 seconds as proof of liveness.
170 logger = lsst.utils.logging.PeriodicLogger(self.task.log, interval=10.0)
171 n_batches = math.ceil(len(inputs) / batch_size)
173 # Loop over the batches
174 for i, batch in enumerate(batches):
175 logger.log("%s/%s batches have been scored.", i, n_batches)
176 torchBlob, labelsList = self.prepare_input(batch)
178 # Run the model
179 with torch.no_grad():
180 output = self.model(torchBlob)
182 # And append the results to the list
183 if i == 0:
184 scores = output
185 else:
186 scores = torch.cat((scores, output.cpu()), dim=0)
188 npyScores = scores.detach().numpy().ravel()
189 return npyScores