19842460617

Committed 08 Oct 2025 10:03AM UTC coverage: 81.663%. Remained the same

Build # 19842460617

Build Type

push

github

Committed by

web-flow

Commit Message

Update README.md, adding reference to GIANT project

Run Details

2053 of 2514 relevant lines covered (81.66%)

1.63 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

58.97

/deepsensor/train/train.py

import deepsensor
from deepsensor.data.task import Task, concat_tasks
from deepsensor.model.convnp import ConvNP

import numpy as np

import lab as B

from typing import List


def set_gpu_default_device() -> None:
    """Set default GPU device for the backend.

    Raises:
        RuntimeError
            If no GPU is available.
        RuntimeError
            If backend is not supported.
        NotImplementedError
            If backend is not supported.

    Returns:
        None.
    """
    if deepsensor.backend.str == "torch":
        # Run on GPU if available
        import torch

        if torch.cuda.is_available():
            # Set default GPU device
            torch.set_default_device("cuda")
            B.set_global_device("cuda:0")
        else:
            raise RuntimeError("No GPU available: torch.cuda.is_available() == False")
    elif deepsensor.backend.str == "tf":
        # Run on GPU if available
        import tensorflow as tf

        if tf.test.is_gpu_available():
            # Set default GPU device
            tf.config.set_visible_devices(
                tf.config.list_physical_devices("GPU")[0], "GPU"
            )
            B.set_global_device("GPU:0")
        else:
            raise RuntimeError("No GPU available: tf.test.is_gpu_available() == False")

    else:
        raise NotImplementedError(f"Backend {deepsensor.backend.str} not implemented")


def train_epoch(
    model: ConvNP,
    tasks: List[Task],
    lr: float = 5e-5,
    batch_size: int = None,
    opt=None,
    progress_bar=False,
    tqdm_notebook=False,
) -> List[float]:
    """Train model for one epoch.

    Args:
        model (:class:`~.model.convnp.ConvNP`):
            Model to train.
        tasks (List[:class:`~.data.task.Task`]):
            List of tasks to train on.
        lr (float, optional):
            Learning rate, by default 5e-5.
        batch_size (int, optional):
            Batch size. Defaults to None. If None, no batching is performed.
        opt (Optimizer, optional):
            TF or Torch optimizer. Defaults to None. If None,
            :class:`tensorflow:tensorflow.keras.optimizer.Adam` is used.
        progress_bar (bool, optional):
            Whether to display a progress bar. Defaults to False.
        tqdm_notebook (bool, optional):
            Whether to use a notebook progress bar. Defaults to False.

    Returns:
        List[float]: List of losses for each task/batch.
    """
    if deepsensor.backend.str == "tf":
        import tensorflow as tf

        if opt is None:
            opt = tf.keras.optimizers.Adam(lr)

        def train_step(tasks):
            if not isinstance(tasks, list):
                tasks = [tasks]
            with tf.GradientTape() as tape:
                task_losses = []
                for task in tasks:
                    task_losses.append(model.loss_fn(task, normalise=True))
                mean_batch_loss = B.mean(B.stack(*task_losses))
            grads = tape.gradient(mean_batch_loss, model.model.trainable_weights)
            opt.apply_gradients(zip(grads, model.model.trainable_weights))
            return mean_batch_loss

    elif deepsensor.backend.str == "torch":
        import torch.optim as optim

        if opt is None:
            opt = optim.Adam(model.model.parameters(), lr=lr)

        def train_step(tasks):
            if not isinstance(tasks, list):
                tasks = [tasks]
            opt.zero_grad()
            task_losses = []
            for task in tasks:
                task_losses.append(model.loss_fn(task, normalise=True))
            mean_batch_loss = B.mean(B.stack(*task_losses))
            mean_batch_loss.backward()
            opt.step()
            return mean_batch_loss.detach().cpu().numpy()

    else:
        raise NotImplementedError(f"Backend {deepsensor.backend.str} not implemented")

    tasks = np.random.permutation(tasks)

    if batch_size is not None:
        n_batches = len(tasks) // batch_size  # Note that this will drop the remainder
    else:
        n_batches = len(tasks)

    if tqdm_notebook:
        from tqdm.notebook import tqdm
    else:
        from tqdm import tqdm

    batch_losses = []
    for batch_i in tqdm(range(n_batches), disable=not progress_bar):
        if batch_size is not None:
            task = concat_tasks(
                tasks[batch_i * batch_size : (batch_i + 1) * batch_size]
            )
        else:
            task = tasks[batch_i]
        batch_loss = train_step(task)
        batch_losses.append(batch_loss)

    return batch_losses


class Trainer:
    """Class for training ConvNP models with an Adam optimiser.

    Args:
        lr (float): Learning rate
    """

    def __init__(self, model: ConvNP, lr: float = 5e-5):
        if deepsensor.backend.str == "tf":
            import tensorflow as tf

            self.opt = tf.keras.optimizers.Adam(lr)
        elif deepsensor.backend.str == "torch":
            import torch.optim as optim

            self.opt = optim.Adam(model.model.parameters(), lr=lr)

        self.model = model

    def __call__(
        self,
        tasks: List[Task],
        batch_size: int = None,
        progress_bar=False,
        tqdm_notebook=False,
    ) -> List[float]:
        """Train model for one epoch."""
        return train_epoch(
            model=self.model,
            tasks=tasks,
            batch_size=batch_size,
            opt=self.opt,
            progress_bar=progress_bar,
            tqdm_notebook=tqdm_notebook,
        )

1	import deepsensor	2✔
2	from deepsensor.data.task import Task, concat_tasks	2✔
3	from deepsensor.model.convnp import ConvNP	2✔
4
5	import numpy as np	2✔
6
7	import lab as B	2✔
8
9	from typing import List	2✔
10
11
12	def set_gpu_default_device() -> None:	2✔
13	"""Set default GPU device for the backend.
14
15	Raises:
16	RuntimeError
17	If no GPU is available.
18	RuntimeError
19	If backend is not supported.
20	NotImplementedError
21	If backend is not supported.
22
23	Returns:
24	None.
25	"""
26	if deepsensor.backend.str == "torch":	×
27	# Run on GPU if available
28	import torch	×
29
30	if torch.cuda.is_available():	×
31	# Set default GPU device
32	torch.set_default_device("cuda")	×
33	B.set_global_device("cuda:0")	×
34	else:
35	raise RuntimeError("No GPU available: torch.cuda.is_available() == False")	×
36	elif deepsensor.backend.str == "tf":	×
37	# Run on GPU if available
38	import tensorflow as tf	×
39
40	if tf.test.is_gpu_available():	×
41	# Set default GPU device
42	tf.config.set_visible_devices(	×
43	tf.config.list_physical_devices("GPU")[0], "GPU"
44	)
45	B.set_global_device("GPU:0")	×
46	else:
47	raise RuntimeError("No GPU available: tf.test.is_gpu_available() == False")	×
48
49	else:
50	raise NotImplementedError(f"Backend {deepsensor.backend.str} not implemented")	×
51
52
53	def train_epoch(	2✔
54	model: ConvNP,
55	tasks: List[Task],
56	lr: float = 5e-5,
57	batch_size: int = None,
58	opt=None,
59	progress_bar=False,
60	tqdm_notebook=False,
61	) -> List[float]:
62	"""Train model for one epoch.
63
64	Args:
65	model (:class:`~.model.convnp.ConvNP`):
66	Model to train.
67	tasks (List[:class:`~.data.task.Task`]):
68	List of tasks to train on.
69	lr (float, optional):
70	Learning rate, by default 5e-5.
71	batch_size (int, optional):
72	Batch size. Defaults to None. If None, no batching is performed.
73	opt (Optimizer, optional):
74	TF or Torch optimizer. Defaults to None. If None,
75	:class:`tensorflow:tensorflow.keras.optimizer.Adam` is used.
76	progress_bar (bool, optional):
77	Whether to display a progress bar. Defaults to False.
78	tqdm_notebook (bool, optional):
79	Whether to use a notebook progress bar. Defaults to False.
80
81	Returns:
82	List[float]: List of losses for each task/batch.
83	"""
84	if deepsensor.backend.str == "tf":	2✔
85	import tensorflow as tf	×
86
87	if opt is None:	×
88	opt = tf.keras.optimizers.Adam(lr)	×
89
90	def train_step(tasks):	×
91	if not isinstance(tasks, list):	×
92	tasks = [tasks]	×
93	with tf.GradientTape() as tape:	×
94	task_losses = []	×
95	for task in tasks:	×
96	task_losses.append(model.loss_fn(task, normalise=True))	×
97	mean_batch_loss = B.mean(B.stack(*task_losses))	×
98	grads = tape.gradient(mean_batch_loss, model.model.trainable_weights)	×
99	opt.apply_gradients(zip(grads, model.model.trainable_weights))	×
100	return mean_batch_loss	×
101
102	elif deepsensor.backend.str == "torch":	2✔
103	import torch.optim as optim	2✔
104
105	if opt is None:	2✔
106	opt = optim.Adam(model.model.parameters(), lr=lr)	×
107
108	def train_step(tasks):	2✔
109	if not isinstance(tasks, list):	2✔
110	tasks = [tasks]	2✔
111	opt.zero_grad()	2✔
112	task_losses = []	2✔
113	for task in tasks:	2✔
114	task_losses.append(model.loss_fn(task, normalise=True))	2✔
115	mean_batch_loss = B.mean(B.stack(*task_losses))	2✔
116	mean_batch_loss.backward()	2✔
117	opt.step()	2✔
118	return mean_batch_loss.detach().cpu().numpy()	2✔
119
120	else:
121	raise NotImplementedError(f"Backend {deepsensor.backend.str} not implemented")	×
122
123	tasks = np.random.permutation(tasks)	2✔
124
125	if batch_size is not None:	2✔
126	n_batches = len(tasks) // batch_size # Note that this will drop the remainder	2✔
127	else:
128	n_batches = len(tasks)	2✔
129
130	if tqdm_notebook:	2✔
131	from tqdm.notebook import tqdm	×
132	else:
133	from tqdm import tqdm	2✔
134
135	batch_losses = []	2✔
136	for batch_i in tqdm(range(n_batches), disable=not progress_bar):	2✔
137	if batch_size is not None:	2✔
138	task = concat_tasks(	2✔
139	tasks[batch_i * batch_size : (batch_i + 1) * batch_size]
140	)
141	else:
142	task = tasks[batch_i]	2✔
143	batch_loss = train_step(task)	2✔
144	batch_losses.append(batch_loss)	2✔
145
146	return batch_losses	2✔
147
148
149	class Trainer:	2✔
150	"""Class for training ConvNP models with an Adam optimiser.
151
152	Args:
153	lr (float): Learning rate
154	"""
155
156	def __init__(self, model: ConvNP, lr: float = 5e-5):	2✔
157	if deepsensor.backend.str == "tf":	2✔
158	import tensorflow as tf	×
159
160	self.opt = tf.keras.optimizers.Adam(lr)	×
161	elif deepsensor.backend.str == "torch":	2✔
162	import torch.optim as optim	2✔
163
164	self.opt = optim.Adam(model.model.parameters(), lr=lr)	2✔
165
166	self.model = model	2✔
167
168	def __call__(	2✔
169	self,
170	tasks: List[Task],
171	batch_size: int = None,
172	progress_bar=False,
173	tqdm_notebook=False,
174	) -> List[float]:
175	"""Train model for one epoch."""
176	return train_epoch(	2✔
177	model=self.model,
178	tasks=tasks,
179	batch_size=batch_size,
180	opt=self.opt,
181	progress_bar=progress_bar,
182	tqdm_notebook=tqdm_notebook,
183	)

alan-turing-institute / deepsensor / 19842460617

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous