diff --git a/pyproject.toml b/pyproject.toml
index b1ef3fd7a835119b9946bcda663f9338216672f6..63190e6f33740abb6cc0ca42c9669bf16aece0e6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "MaggotUBA-adapter"
-version = "0.7"
+version = "0.8"
description = "Interface between MaggotUBA and the Nyx tagging UI"
authors = ["François Laurent"]
license = "MIT"
diff --git a/src/maggotuba/models/modules.py b/src/maggotuba/models/modules.py
index fbb3b844b6bd4212872ff016343acfded091ed74..9d197a3885a6719b3003ce793b84e5d0e4fa4a2f 100644
--- a/src/maggotuba/models/modules.py
+++ b/src/maggotuba/models/modules.py
@@ -90,6 +90,9 @@ class MaggotModule(nn.Module):
def parameters(self, recurse=True):
return self.model.parameters(recurse)
+ def to(self, device):
+ self.model.to(device)
+
"""
Initialize a model's weights and bias (if any).
@@ -305,6 +308,9 @@ class DeepLinear(nn.Module):
torch.save(self.state_dict(), path)
check_permissions(path)
+ def to(self, device):
+ self.layers.to(device)
+
class MaggotClassifier(MaggotModule):
def __init__(self, path, behavior_labels=[], n_latent_features=None,
n_layers=1, cfgfile=None, ptfile="trained_classifier.pt"):
@@ -374,6 +380,9 @@ class SupervisedMaggot(nn.Module):
def forward(self, x):
return self.clf(self.encoder(x))
+ def mask_forward(self, x):
+ return self.clf(self.encoder.mask_forward(x))
+
def save(self):
enc, clf = self.encoder, self.clf
enc.save()
@@ -385,6 +394,10 @@ class SupervisedMaggot(nn.Module):
self.clf.model # force parameter loading or initialization
return super().parameters(self)
+ def to(self, device):
+ self.encoder.to(device)
+ self.clf.to(device)
+
class MultiscaleSupervisedMaggot(nn.Module):
def __init__(self, cfgfilepath, behaviors=[], n_layers=1):
super().__init__()
@@ -413,3 +426,45 @@ class MultiscaleSupervisedMaggot(nn.Module):
self.clf.model # force parameter loading or initialization
return super().parameters(self)
+"""
+Bagging for `SupervisedMaggot`.
+
+Taggers in the bag are individually trained.
+For now the bag itself cannot be trained and is used for prediction only.
+
+Bags of taggers are stored so that the models directory only contains
+subdirectories, each subdirectory specifying an individual tagger.
+"""
+class MaggotBag(nn.Module):
+ def __init__(self, paths, behaviors=[], n_layers=1, cls=SupervisedMaggot):
+ super().__init__()
+ self.maggots = [cls(path, behaviors, n_layers) for path in paths]
+ self._lead_maggot = None
+
+ def forward(self, x):
+ #return torch.cat([encoder.mask_forward(x) for encoder in self.encoders], dim=1)
+ return self.vote([maggot.mask_forward(x) for maggot in self.maggots])
+
+ def vote(self, y):
+ vote, _ = torch.mode(torch.stack(y, dim=len(y[0].shape)))
+ return vote
+
+ @property
+ def encoder(self):
+ return self.maggots[self.lead_maggot].encoder
+
+ @property
+ def clf(self):
+ return self.maggots[self.lead_maggot].clf
+
+ @property
+ def lead_maggot(self):
+ if self._lead_maggot is None:
+ len_traj = 0
+ for i, maggot in enumerate(self.maggots):
+ len_traj_ = maggot.encoder.config['len_traj']
+ if len_traj < len_traj_:
+ len_traj = len_traj_
+ self._lead_maggot = i
+ return self._lead_maggot
+
diff --git a/src/maggotuba/models/predict_model.py b/src/maggotuba/models/predict_model.py
index aa320d8546320dc8096d806b8b7bc5aef43a4fc7..c1abd548e048cd4b73d07b6d9c9f2a1f45a15da0 100644
--- a/src/maggotuba/models/predict_model.py
+++ b/src/maggotuba/models/predict_model.py
@@ -1,6 +1,6 @@
from taggingbackends.data.labels import Labels
from taggingbackends.features.skeleton import get_5point_spines
-from maggotuba.models.trainers import MaggotTrainer, MultiscaleMaggotTrainer, new_generator
+from maggotuba.models.trainers import MaggotTrainer, MultiscaleMaggotTrainer, MaggotBagging, new_generator
import numpy as np
import logging
@@ -27,20 +27,25 @@ def predict_model(backend, **kwargs):
assert 0 < len(input_files_and_labels)
# load the model
model_files = backend.list_model_files()
- config_file = [file for file in model_files if file.name.endswith("config.json")]
- n_config_files = len(config_file)
- if n_config_files == 0:
- raise RuntimeError(f"no such tagger found: {backend.model_instance}")
- config_file = [file
- for file in config_file
- if file.name.endswith("clf_config.json")
- and file.parent == backend.model_dir()]
- assert len(config_file) == 1
- config_file = config_file[-1]
- if 2 < n_config_files:
- model = MultiscaleMaggotTrainer(config_file)
+ config_files = [file
+ for file in model_files
+ if file.name.endswith('config.json')]
+ if len(config_files) == 0:
+ raise RuntimeError(f"no config files found for tagger: {backend.model_instance}")
+ single_encoder_classifier = len(config_files) == 2
+ config_files = [file
+ for file in config_files
+ if file.name == 'clf_config.json']
+ if len(config_files) == 0:
+ raise RuntimeError(f"no classifier config files found; is {backend.model_instance} tagger trained?")
+ elif len(config_files) == 1:
+ config_file = config_files[0]
+ if single_encoder_classifier:
+ model = MaggotTrainer(config_file)
+ else:
+ model = MultiscaleMaggotTrainer(config_file)
else:
- model = MaggotTrainer(config_file)
+ model = MaggotBagging(config_files)
#
if len(input_files) == 1:
input_files = next(iter(input_files.values()))
@@ -133,9 +138,9 @@ def predict_individual_data_files(backend, model, input_files_and_labels):
#
done = True
-def predict_larva_dataset(backend, model, file, subset="validation"):
+def predict_larva_dataset(backend, model, file, subset="validation", subsets=(.8, .2, 0)):
from taggingbackends.data.dataset import LarvaDataset
- dataset = LarvaDataset(file, new_generator())
+ dataset = LarvaDataset(file, new_generator(), subsets)
return model.predict(dataset, subset)
def _zip(xs, ys):
diff --git a/src/maggotuba/models/trainers.py b/src/maggotuba/models/trainers.py
index 025b95dbbb83af436d92fd9d53f6b450d92c1ab6..aba17748fb6256dcd9c453da168b45d0e0b5a2d7 100644
--- a/src/maggotuba/models/trainers.py
+++ b/src/maggotuba/models/trainers.py
@@ -3,7 +3,7 @@ import torch
import torch.nn as nn
from behavior_model.models.neural_nets import device
#import behavior_model.data.utils as data_utils
-from maggotuba.models.modules import SupervisedMaggot, MultiscaleSupervisedMaggot
+from maggotuba.models.modules import SupervisedMaggot, MultiscaleSupervisedMaggot, MaggotBag
from taggingbackends.features.skeleton import interpolate
"""
@@ -231,7 +231,7 @@ class MaggotTrainer:
self.model.save()
def new_generator(seed=None):
- generator = torch.Generator(device)
+ generator = torch.Generator('cpu')
if seed == 'random': return generator
if seed is None: seed = 0b11010111001001101001110
return generator.manual_seed(seed)
@@ -265,6 +265,14 @@ class MultiscaleMaggotTrainer(MaggotTrainer):
return self._default_encoder_config
+class MaggotBagging(MaggotTrainer):
+ def __init__(self, cfgfilepaths, behaviors=[], n_layers=1,
+ average_body_length=1.0, device=device):
+ self.model = MaggotBag(cfgfilepaths, behaviors, n_layers)
+ self.average_body_length = average_body_length # usually set later
+ self.device = device
+
+
"""
Pick the adequate trainer following a rapid inspection of the config file(s).
@@ -272,6 +280,8 @@ For now, config files are actually not inspected. However, using this function
is highly recommended as more models are introduced with future releases.
"""
def make_trainer(config_file, *args, **kwargs):
+ # the type criterion does not fail in the case of unimplemented bagging,
+ # as config files are listed in a pretrained_models subdirectory.
if isinstance(config_file, list): # multiple encoders
config_files = config_file
model = MultiscaleMaggotTrainer(config_files, *args, **kwargs)