DLoc baseline + Mamba minimal run

.gitignore

@@ -0,0 +1,12 @@
+# Data and outputs
+/data/
+/runs/
+**/*.mat
+**/*.h5
+**/*.hdf5
+**/*.pt
+**/*.pth
+**/*.ckpt
+**/*.npz
+**/*.npy
+**/*.log

Generators.py

@@ -0,0 +1,145 @@
+#!/usr/bin/python
+'''
+Defines the Encoders and Decoder that make up the
+Enc+Dec/Enc+2Dec model defined in DLoc
+Idea and code originally from fstin Johnson's architecture.
+https://github.com/jcjohnson/fast-neural-style/
+Code base expanded from pix2pix Phillip Isola's implementation
+https://github.com/phillipi/pix2pix
+
+You can add your costum network building blocks here to test various other architectures.
+'''
+import torch
+import torch.nn as nn
+import functools
+
+
+# Tha base Encoder function defined for the DLoc's Encoder
+class ResnetEncoder(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='zero'):
+        assert(n_blocks >= 0)
+        super(ResnetEncoder, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.Conv2d(input_nc, input_nc, kernel_size=7, padding=3,
+                           bias=use_bias),
+                 norm_layer(input_nc),
+                 nn.Tanh()]
+
+        model += [nn.Conv2d(input_nc, ngf, kernel_size=7, padding=3,
+                           bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,
+                                stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2),
+                      nn.ReLU(True)]
+
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        return self.model(input)
+
+
+# Tha base Decoder function defined for the DLoc's Decoders.
+# Depending upon the ModelADT wraper around the decoder,
+# the decoder would either be a Location Decoder or a Consistency decoder.
+# For more details refer to ModelADT.py wrapper implementation and params.py
+class ResnetDecoder(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=9, padding_type='zero', encoder_blocks=6):
+        assert(n_blocks >= 0)
+        super(ResnetDecoder, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = []
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            if i <= encoder_blocks:
+                continue
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                         kernel_size=3, stride=2,
+                                         padding=1, output_padding=0,
+                                         bias=use_bias),
+                      norm_layer(int(ngf * mult / 2)),
+                      nn.ReLU(True)]
+#         model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=[3,3])]
+        model += [nn.Sigmoid()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        #print("decoder.input = ", input.shape)
+        return self.model(input)
+# In[3]:
+
+
+# Define a resnet block
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
+                       norm_layer(dim),
+                       nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.25)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
+                       norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        out = x + self.conv_block(x)
+        return out

README.md

@@ -0,0 +1,37 @@
+# DLoc minimal run (baseline + mamba)
+
+This folder is an isolated, minimal run package for:
+- DLoc baseline (Fig 10b, single GPU)
+- Mamba-DLoc (single GPU)
+
+It excludes all data and training outputs by default.
+
+## Required data (place in ./data)
+- dataset_jacobs_July28.mat (fixed 7,635 samples)
+- dataset_non_fov_train_jacobs_July28_2.mat
+- dataset_fov_train_jacobs_July28_2.mat
+- dataset_non_fov_test_jacobs_July28_2.mat
+- dataset_fov_test_jacobs_July28_2.mat
+
+## Setup
+This assumes PyTorch is already installed (e.g., vast.ai PyTorch image).
+
+```bash
+pip install -r requirements.txt
+```
+
+## Run baseline (Fig 10b)
+```bash
+cp params_storage/params_fig10b_single_gpu.py params.py
+python train_and_test.py
+```
+
+## Run Mamba
+```bash
+cp params_storage/params_mamba_single_gpu.py params.py
+python train_and_test.py
+```
+
+## Notes
+- Outputs go to ./runs/<timestamp>/
+- This folder is safe to upload to Git/HF (data and runs are ignored).

data_loader.py

@@ -0,0 +1,98 @@
+#!/usr/bin/python
+# A simple data loader that imports the train and test mat files
+# from the `filename` and converts them to torch.tesnors()
+# to be loaded for training and testing DLoc network
+# `features_wo_offset`: targets for the consistency decoder
+# `features_w_offset` : inputs for the network/encoder
+# `labels_gaussian_2d`: targets for the location decoder
+import torch
+import h5py
+import scipy.io
+import numpy as np
+
+def load_data(filename):
+    print('Loading '+filename)
+    arrays = {}
+    f = h5py.File(filename,'r')
+    features_wo_offset = torch.tensor(np.transpose(np.array(f.get('features_wo_offset'), dtype=np.float32)), dtype=torch.float32)
+    features_w_offset = torch.tensor(np.transpose(np.array(f.get('features_w_offset'), dtype=np.float32)), dtype=torch.float32)
+    labels_gaussian_2d = torch.tensor(np.transpose(np.array(f.get('labels_gaussian_2d'), dtype=np.float32)), dtype=torch.float32)
+        
+    return features_wo_offset,features_w_offset, labels_gaussian_2d
+
+
+class LazyMultiHDF5Dataset(torch.utils.data.Dataset):
+    """Reads samples on demand from multiple HDF5 .mat files.
+    Uses ~0 RAM regardless of dataset size. Each __getitem__ does 3 small
+    HDF5 slice reads (~700 KB total per sample).
+
+    Returns the same tuple ordering as the original TensorDataset:
+        (features_wo_offset, features_w_offset, labels_gaussian_2d)
+    """
+
+    def __init__(self, file_paths):
+        self.file_paths = file_paths
+        self._handles = {}           # lazy-opened per worker/process
+        self.cum_sizes = []
+        self.var_names = {}          # per file: (w_name, wo_name, lbl_name)
+
+        total = 0
+        for fp in file_paths:
+            with h5py.File(fp, 'r') as f:
+                keys = list(f.keys())
+                w_name  = 'features_w_offset'  if 'features_w_offset'  in keys else 'features_with_offset'
+                wo_name = 'features_wo_offset' if 'features_wo_offset' in keys else 'features_without_offset'
+                lbl_name = 'labels_gaussian_2d'
+                self.var_names[fp] = (w_name, wo_name, lbl_name)
+                # HDF5 stores MATLAB dims reversed: last dim = N
+                n = f[w_name].shape[-1]
+                total += n
+                self.cum_sizes.append(total)
+                print(f'  Indexed {fp}: {n} samples')
+
+        self.total_len = total
+        print(f'  Total: {self.total_len} samples (lazy, ~0 MB RAM)')
+
+    def __len__(self):
+        return self.total_len
+
+    def _get_handle(self, fp):
+        if fp not in self._handles:
+            self._handles[fp] = h5py.File(fp, 'r')
+        return self._handles[fp]
+
+    def __getitem__(self, idx):
+        # Find which file this idx belongs to
+        file_idx = 0
+        for i, cs in enumerate(self.cum_sizes):
+            if idx < cs:
+                file_idx = i
+                break
+
+        offset = self.cum_sizes[file_idx - 1] if file_idx > 0 else 0
+        local_idx = idx - offset
+
+        fp = self.file_paths[file_idx]
+        f = self._get_handle(fp)
+        w_name, wo_name, lbl_name = self.var_names[fp]
+
+        # HDF5 dims are reversed vs MATLAB:
+        #   features: MATLAB [N,4,161,361] -> HDF5 [361,161,4,N]
+        #   labels:   MATLAB [N,161,361]   -> HDF5 [361,161,N]
+        feat_w  = np.array(f[w_name][:, :, :, local_idx],  dtype=np.float32).transpose()  # -> [4,161,361]
+        feat_wo = np.array(f[wo_name][:, :, :, local_idx], dtype=np.float32).transpose()  # -> [4,161,361]
+        label   = np.array(f[lbl_name][:, :, local_idx],   dtype=np.float32).transpose()  # -> [161,361]
+        label   = label[np.newaxis, :, :]   # -> [1,161,361]
+
+        return (
+            torch.from_numpy(feat_wo),
+            torch.from_numpy(feat_w),
+            torch.from_numpy(label)
+        )
+
+    def __del__(self):
+        for f in self._handles.values():
+            try:
+                f.close()
+            except Exception:
+                pass

joint_model.py

@@ -0,0 +1,217 @@
+#!/usr/bin/python
+'''
+Defines a PyTorch graph for forward and backward propogation
+within the network encoders and decoders.
+'''
+import torch
+from utils import *
+from Generators import *
+from data_loader import *
+from params import *
+
+
+'''
+Class Definition for 1 Encoder and 1 Decoder joint model
+'''
+class Enc_Dec_Network():
+
+    def initialize(self, opt, encoder, decoder, frozen_dec=False, frozen_enc=False, gpu_ids='1'):
+        self.opt = opt
+        self.isTrain = opt.isTrain
+        self.encoder = encoder
+        self.decoder = decoder
+        self.frozen_dec = frozen_dec
+        self.frozen_enc = frozen_enc
+        self.device = torch.device('cuda:{}'.format(gpu_ids[0])) # if self.gpu_ids else torch.device('cpu')
+        # self.encoder.net = encoder.net.to(self.device)
+        # self.decoder.net = decoder.net.to(self.device)
+
+    def set_input(self, input, target, convert_enc=True, shuffle_channel=True):
+        self.input = input.to(self.device)
+        self.target = target.to(self.device)
+        self.encoder.set_data(self.input, self.input, convert=convert_enc, shuffle_channel=shuffle_channel)
+    
+    def save_networks(self, epoch):
+        self.encoder.save_networks(epoch)
+        self.decoder.save_networks(epoch)
+
+    def save_outputs(self):
+        self.encoder.save_outputs()
+        self.decoder.save_outputs()
+    
+    def update_learning_rate(self):
+        self.encoder.update_learning_rate()
+        
+    def forward(self):
+        self.encoder.forward()
+        self.decoder.set_data(self.encoder.output, self.target)
+        self.decoder.forward()
+    
+    def test(self):
+        self.encoder.test()
+        self.decoder.set_data(self.encoder.output, self.target)
+        self.decoder.test()
+
+    def backward(self):
+        self.decoder.backward()
+        # self.encoder.backward()
+        
+    def optimize_parameters(self):
+        self.forward()
+        self.backward()
+        if not self.frozen_enc:
+            self.encoder.optimizer.step()
+        if not self.frozen_dec:
+            self.decoder.optimizer.step()
+        self.encoder.optimizer.zero_grad()
+        self.decoder.optimizer.zero_grad()
+
+    def eval(self):
+        self.encoder.forward()
+        self.decoder.set_data(self.encoder.output, self.target)
+        self.decoder.forward()
+
+'''
+Class Definition for the final DLoc architecture with
+1 Encoder and 2 Decoders joint model
+'''
+class Enc_2Dec_Network():
+
+    def initialize(self, opt , encoder, decoder, offset_decoder, frozen_dec=False, frozen_enc=False, gpu_ids='1'):
+        print('initializing Encoder and 2 Decoders Model')
+        self.opt = opt
+        self.isTrain = opt.isTrain      
+        self.encoder = encoder
+        self.decoder = decoder
+        self.offset_decoder = offset_decoder
+        self.frozen_dec = frozen_dec
+        self.frozen_enc = frozen_enc
+        self.device = torch.device('cuda:{}'.format(gpu_ids[0])) # if self.gpu_ids else torch.device('cpu')
+        # self.encoder.net = encoder.net.to(self.device)
+        # self.decoder.net = decoder.net.to(self.device)
+        self.results_save_dir = opt.results_dir
+
+    def set_input(self, input, target ,offset_target ,convert_enc=True, shuffle_channel=True):
+        # features_w_offset, labels_gaussian_2d, features_wo_offset
+        # input,             target,             offset_target
+        self.input = input.to(self.device)
+        self.target = target.to(self.device)
+        self.offset_target = offset_target.to(self.device)      
+        self.encoder.set_data(self.input, self.input, convert=convert_enc, shuffle_channel=shuffle_channel)
+    
+    def save_networks(self, epoch):
+        self.encoder.save_networks(epoch)
+        self.decoder.save_networks(epoch)
+        self.offset_decoder.save_networks(epoch)
+
+    def save_outputs(self):
+        self.encoder.save_outputs()
+        self.decoder.save_outputs()
+        self.offset_decoder.save_outputs()
+    
+    def update_learning_rate(self):
+        self.encoder.update_learning_rate()
+        self.decoder.update_learning_rate()
+        self.offset_decoder.update_learning_rate()
+        
+    def forward(self):
+        self.encoder.forward()
+        self.decoder.set_data(self.encoder.output, self.target)
+        self.offset_decoder.set_data(self.encoder.output, self.offset_target)
+        self.decoder.forward()
+        self.offset_decoder.forward()
+    
+    # Test the network once set into Evaluation mode!
+    def test(self):      
+        self.encoder.test()
+        self.decoder.set_data(self.encoder.output, self.target)
+        self.offset_decoder.set_data(self.encoder.output, self.offset_target)
+        self.decoder.test()
+        self.offset_decoder.test()
+
+    def backward(self):
+        self.decoder.backward()
+        self.offset_decoder.backward()
+        # self.encoder.backward()
+        
+    def optimize_parameters(self):
+        self.forward()
+        self.backward()
+        if not self.frozen_enc:
+            self.encoder.optimizer.step()
+        if not self.frozen_dec:
+            self.decoder.optimizer.step()
+            self.offset_decoder.optimizer.step()
+        self.encoder.optimizer.zero_grad()
+        self.decoder.optimizer.zero_grad()
+        self.offset_decoder.optimizer.zero_grad()
+
+    # set the models to evaluation mode
+    def eval(self):
+        self.encoder.eval()
+        self.decoder.eval()
+        self.offset_decoder.eval()
+
+
+class _EncoderProxy:
+    """Lightweight proxy that shares opt/optimizer but no-ops on LR update.
+
+    trainer.py calls update_learning_rate() on both model.encoder and
+    model.decoder.  For the Mamba pipeline there is only one optimizer,
+    so this proxy prevents double-stepping the LR scheduler.
+    """
+
+    def __init__(self, model_adt):
+        self.opt = model_adt.opt
+        self.optimizer = model_adt.optimizer
+
+    def update_learning_rate(self):
+        pass  # decoder handles the single scheduler
+
+    def save_networks(self, epoch):
+        pass  # decoder handles saving
+
+
+class Mamba_Network():
+    """Joint-model wrapper for the single end-to-end MambaDLocNet.
+
+    Exposes the same interface that trainer.py expects
+    (model.encoder, model.decoder, set_input, optimize_parameters, ...)
+    while internally using a single ModelADT with one optimizer.
+    """
+
+    def initialize(self, opt, mamba_model_adt, gpu_ids='0'):
+        self.opt = opt
+        self.isTrain = opt.isTrain
+        self.decoder = mamba_model_adt           # trainer reads .output, .loss, .model_name, .opt
+        self.encoder = _EncoderProxy(mamba_model_adt)  # no-op LR update
+        self.device = torch.device('cuda:{}'.format(gpu_ids[0]))
+
+    def set_input(self, input, target, shuffle_channel=False):
+        self.input = input.to(self.device)
+        self.target = target.to(self.device)
+        self.decoder.set_data(self.input, self.target, shuffle_channel=shuffle_channel)
+
+    def save_networks(self, epoch):
+        self.decoder.save_networks(epoch)
+
+    def update_learning_rate(self):
+        self.decoder.update_learning_rate()
+
+    def forward(self):
+        self.decoder.forward()
+
+    def test(self):
+        self.decoder.test()
+
+    def backward(self):
+        self.decoder.backward()
+
+    def optimize_parameters(self):
+        self.forward()
+        self.backward()
+        self.decoder.optimizer.step()
+        self.decoder.optimizer.zero_grad()
+
+    def eval(self):
+        self.decoder.eval()

mamba_model.py

@@ -0,0 +1,273 @@
+#!/usr/bin/env python
+"""
+mamba_model.py
+Mamba-DLoc: Replaces the ResNet bottleneck in DLoc with Mamba selective
+state-space (S6) blocks for efficient long-range sequence modelling.
+
+Architecture overview
+---------------------
+CNN Stem  (4 -> 128 channels, 2x spatial downsample)
+  -> Tokenize (flatten spatial -> sequence)
+  -> 4x MambaResidualBlock  (LayerNorm -> Mamba -> residual)
+  -> Reshape back to spatial
+  -> CNN Decoder (upsample back to original resolution)
+
+Input:  [B, 4, 161, 361]  AoA heatmaps
+Output: [B, 1, 161, 361]  location heatmap
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Try to import the CUDA-optimised Mamba implementation (Linux + CUDA only).
+# Fall back to a pure-PyTorch implementation on Windows / CPU.
+# ---------------------------------------------------------------------------
+HAS_MAMBA_SSM = False
+try:
+    from mamba_ssm import Mamba as CudaMamba  # type: ignore
+    HAS_MAMBA_SSM = True
+except ImportError:
+    CudaMamba = None
+
+
+# ===================================================================
+# Pure-PyTorch Mamba (S6 selective scan) - works on any backend
+# ===================================================================
+class PurePytorchMamba(nn.Module):
+    """
+    Minimal re-implementation of the Mamba selective-scan block (S6).
+
+    Interface: [B, L, D] -> [B, L, D]
+
+    Parameters
+    ----------
+    d_model : int   - input / output dimension (D)
+    d_state : int   - SSM hidden state dimension (N)
+    d_conv  : int   - local convolution kernel width
+    expand  : int   - expansion factor for inner dimension (E = expand * D)
+    """
+
+    def __init__(self, d_model: int = 128, d_state: int = 16,
+                 d_conv: int = 4, expand: int = 2):
+        super().__init__()
+        self.d_model = d_model
+        self.d_state = d_state
+        self.d_conv = d_conv
+        self.expand = expand
+        self.d_inner = d_model * expand           # E
+        self.dt_rank = math.ceil(d_model / 16)     # rank for delta-t projection
+
+        # --- input projection: D -> 2*E  (main path + gate) ---
+        self.in_proj = nn.Linear(d_model, 2 * self.d_inner, bias=False)
+
+        # --- causal depthwise conv on main path ---
+        self.conv1d = nn.Conv1d(
+            self.d_inner, self.d_inner,
+            kernel_size=d_conv, padding=d_conv - 1,  # causal: trim later
+            groups=self.d_inner, bias=True,
+        )
+
+        # --- SSM parameter projections ---
+        #   from the convolved main path (E) -> dt_rank + 2*N
+        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + 2 * d_state, bias=False)
+
+        #   dt_rank -> E  (broadcast delta across inner dim)
+        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
+
+        # --- learnable SSM parameters ---
+        # A_log: [E, N]  initialised as log(1..N) broadcast to each inner dim
+        A = torch.arange(1, d_state + 1, dtype=torch.float32).unsqueeze(0).expand(self.d_inner, -1)
+        self.A_log = nn.Parameter(torch.log(A))           # learned in log-space
+
+        # D (skip connection scalar per inner dim)
+        self.D = nn.Parameter(torch.ones(self.d_inner))
+
+        # --- output projection: E -> D ---
+        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
+
+    # ---------------------------------------------------------------
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x : [B, L, D]
+        returns : [B, L, D]
+        """
+        B, L, _ = x.shape
+
+        # --- input projection & split ---
+        xz = self.in_proj(x)                              # [B, L, 2E]
+        x_main, z = xz.chunk(2, dim=-1)                   # each [B, L, E]
+
+        # --- causal conv1d on main path ---
+        # conv1d expects [B, E, L]
+        x_main = x_main.transpose(1, 2)                   # [B, E, L]
+        x_main = self.conv1d(x_main)[:, :, :L]            # causal trim
+        x_main = x_main.transpose(1, 2)                   # [B, L, E]
+        x_main = F.silu(x_main)
+
+        # --- SSM parameters ---
+        x_ssm = self.x_proj(x_main)                       # [B, L, dt_rank + 2N]
+        dt, B_param, C_param = x_ssm.split(
+            [self.dt_rank, self.d_state, self.d_state], dim=-1
+        )
+        # dt: [B, L, dt_rank] -> [B, L, E]
+        dt = F.softplus(self.dt_proj(dt))                  # [B, L, E]
+
+        # A: [E, N]  (negative for stability)
+        A = -torch.exp(self.A_log)                         # [E, N]
+
+        # --- discretize ---
+        # dA: [B, L, E, N]
+        dA = torch.exp(dt.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))
+        # dB: [B, L, E, N]
+        dB = dt.unsqueeze(-1) * B_param.unsqueeze(2)       # broadcast E
+
+        # --- selective scan (sequential, kept simple for correctness) ---
+        # x_main: [B, L, E]  ->  weight input
+        # h: [B, E, N]       SSM hidden state
+        h = torch.zeros(B, self.d_inner, self.d_state,
+                        device=x.device, dtype=x.dtype)
+        ys = []
+        for t in range(L):
+            h = dA[:, t, :, :] * h + dB[:, t, :, :] * x_main[:, t, :].unsqueeze(-1)
+            y_t = (h * C_param[:, t, :].unsqueeze(1)).sum(dim=-1)   # [B, E]
+            ys.append(y_t)
+        y = torch.stack(ys, dim=1)                         # [B, L, E]
+
+        # --- skip connection ---
+        y = y + x_main * self.D.unsqueeze(0).unsqueeze(0)
+
+        # --- gate ---
+        y = y * F.silu(z)
+
+        # --- output projection ---
+        return self.out_proj(y)                            # [B, L, D]
+
+
+# ===================================================================
+# Mamba residual block: LayerNorm -> Mamba -> residual add
+# ===================================================================
+class MambaResidualBlock(nn.Module):
+    """
+    Pre-norm residual block wrapping a Mamba layer.
+
+    Interface: [B, L, D] -> [B, L, D]
+    Uses the CUDA Mamba kernel when available, otherwise falls back to
+    PurePytorchMamba.
+    """
+
+    def __init__(self, d_model: int = 128, d_state: int = 16,
+                 d_conv: int = 4, expand: int = 2):
+        super().__init__()
+        self.norm = nn.LayerNorm(d_model)
+        if HAS_MAMBA_SSM and CudaMamba is not None:
+            self.mamba = CudaMamba(
+                d_model=d_model, d_state=d_state,
+                d_conv=d_conv, expand=expand,
+            )
+        else:
+            self.mamba = PurePytorchMamba(
+                d_model=d_model, d_state=d_state,
+                d_conv=d_conv, expand=expand,
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.mamba(self.norm(x))
+
+
+# ===================================================================
+# MambaDLocNet - full end-to-end model
+# ===================================================================
+class MambaDLocNet(nn.Module):
+    """
+    Mamba-DLoc: CNN stem  ->  Mamba sequence blocks  ->  CNN decoder.
+
+    Input:  [B, 4, 161, 361]
+    Output: [B, 1, 161, 361]
+    """
+
+    # Spatial sizes at each resolution level (full, mid, low)
+    SPATIAL_SIZES = [(161, 361), (81, 181), (41, 91)]
+
+    def __init__(self, input_nc: int = 4, d_model: int = 128,
+                 n_mamba_blocks: int = 4,
+                 d_state: int = 16, d_conv: int = 4, expand: int = 2):
+        super().__init__()
+        self.d_model = d_model
+
+        # ------ CNN Stem (encoder) ------
+        self.stem = nn.Sequential(
+            # [B,input_nc,161,361] -> [B,32,161,361]
+            nn.Conv2d(input_nc, 32, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            # [B,32,161,361] -> [B,64,81,181]
+            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            # [B,64,81,181] -> [B,128,41,91]
+            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+        )
+
+        # ------ Mamba sequence blocks ------
+        self.mamba_blocks = nn.Sequential(*[
+            MambaResidualBlock(d_model=d_model, d_state=d_state,
+                               d_conv=d_conv, expand=expand)
+            for _ in range(n_mamba_blocks)
+        ])
+
+        # ------ CNN Decoder ------
+        # Stage 1: [B,128,41,91] -> interpolate to (81,181) -> conv 128->64
+        self.dec_conv1 = nn.Sequential(
+            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+        )
+        # Stage 2: [B,64,81,181] -> interpolate to (161,361) -> conv 64->32
+        self.dec_conv2 = nn.Sequential(
+            nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+        )
+        # Stage 3: 1x1 head -> sigmoid
+        self.dec_head = nn.Sequential(
+            nn.Conv2d(32, 1, kernel_size=1, bias=True),
+            nn.Sigmoid(),
+        )
+
+    # ---------------------------------------------------------------
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x : [B, 4, 161, 361]
+        returns : [B, 1, 161, 361]
+        """
+        B = x.size(0)
+
+        # --- CNN Stem ---
+        feat = self.stem(x)                                # [B, 128, 41, 91]
+
+        # --- Tokenize: flatten spatial dims into sequence ---
+        _, C, H, W = feat.shape
+        tokens = feat.reshape(B, C, H * W).permute(0, 2, 1)  # [B, 3731, 128]
+
+        # --- Mamba blocks ---
+        tokens = self.mamba_blocks(tokens)                 # [B, 3731, 128]
+
+        # --- Reshape back to spatial ---
+        feat = tokens.permute(0, 2, 1).reshape(B, C, H, W)   # [B, 128, 41, 91]
+
+        # --- CNN Decoder ---
+        feat = F.interpolate(feat, size=self.SPATIAL_SIZES[1],
+                             mode='bilinear', align_corners=False)   # [B,128,81,181]
+        feat = self.dec_conv1(feat)                                   # [B,64,81,181]
+
+        feat = F.interpolate(feat, size=self.SPATIAL_SIZES[0],
+                             mode='bilinear', align_corners=False)   # [B,64,161,361]
+        feat = self.dec_conv2(feat)                                   # [B,32,161,361]
+
+        out = self.dec_head(feat)                                     # [B,1,161,361]
+        return out

modelADT.py

@@ -0,0 +1,289 @@
+#!/usr/bin/python
+'''
+Defines a generic wrapper class for all the network models
+Utilizes params.py to create, initiate, load and train the network.
+'''
+import torch
+import numpy as np
+from collections import OrderedDict
+import scipy.io
+from torch.autograd import Variable
+import os
+from Generators import *
+from utils import *
+
+class ModelADT():
+    def name(self):
+        return 'Pix2PixModel'
+
+    def initialize(self, opt):
+        
+        self.opt = opt
+#        self.gpu_ids = opt.gpu_ids
+        gpu_ids = []
+        for i in range(torch.cuda.device_count()):
+            gpu_ids.append(str(i))
+        print(gpu_ids)
+        self.gpu_ids = gpu_ids
+        self.isTrain = opt.isTrain
+        self.loss_weight = opt.lambda_L
+        self.reg_loss_weight = opt.lambda_reg
+        self.cross_loss_weight = opt.lambda_cross
+        self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) # if self.gpu_ids else torch.device('cpu')
+        
+        print(self.device)
+        self.model_name = self.opt.name
+        self.save_dir = os.path.join(self.opt.checkpoints_save_dir, self.model_name)
+        self.load_dir = os.path.join(self.opt.checkpoints_load_dir, self.model_name)
+        self.results_save_dir = opt.results_dir
+
+        self.loss_names = []
+        self.visual_names = []
+        self.image_paths = []
+
+        self.loss_names = ['loss_criterion']
+        self.visual_names = ['output']
+        
+
+        self.net = get_model_funct(self.opt.net)(self.opt, self.gpu_ids)
+        self.net = self.net.to(self.device)
+
+        if self.isTrain:
+            if self.opt.loss_type == "L2":
+                self.loss_criterion = torch.nn.MSELoss()
+            elif self.opt.loss_type == "L1":
+                self.loss_criterion = torch.nn.L1Loss()
+            elif self.opt.loss_type == "L1_sumL2":
+                self.loss_criterion = torch.nn.L1Loss()
+            elif self.opt.loss_type == "L2_sumL2":
+                self.loss_criterion = torch.nn.MSELoss()
+            elif self.opt.loss_type == "L2_sumL1":
+                self.loss_criterion = torch.nn.MSELoss()
+            elif self.opt.loss_type == "L2_offset_loss":
+                self.loss_criterion = torch.nn.MSELoss()
+            elif self.opt.loss_type == "L1_offset_loss":
+                self.loss_criterion = torch.nn.L1Loss()
+            elif self.opt.loss_type == "L1_sumL2_cross":
+                self.loss_criterion = torch.nn.MSELoss()
+            elif self.opt.loss_type == "L2_sumL2_cross":
+                self.loss_criterion = torch.nn.MSELoss()
+                self.cross_loss_criterion = torch.nn.NLLLoss()
+
+            # initialize optimizers
+            self.optimizers = []
+            self.optimizer = torch.optim.Adam(self.net.parameters(),lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=opt.weight_decay)
+            self.optimizers.append(self.optimizer)
+
+    # load and print networks; create schedulers
+    def setup(self, opt, parser=None):
+        if self.isTrain:
+            self.schedulers = [get_scheduler(optimizer, opt) for optimizer in self.optimizers]
+
+        if not self.isTrain or opt.continue_train:
+            self.load_networks(opt.starting_epoch_count)
+            if self.isTrain and opt.continue_train:
+                self.load_training_state(opt.starting_epoch_count)
+        self.print_networks(opt.verbose)
+
+    # make models eval mode during test time
+    def eval(self):
+        self.net.eval()
+
+    # used in test time, wrapping `forward` in no_grad() so we don't save
+    # intermediate steps for backprop
+    def test(self):
+        with torch.no_grad():
+            self.forward()
+    
+    def save_outputs(self):
+        if not os.path.exists(self.results_save_dir):
+            os.makedirs(self.results_save_dir, exist_ok=True)
+        to_save_dict = {}
+        for tensor in self.visual_names:
+            tensor_val = getattr(self, tensor).data.cpu().numpy()
+            to_save_dict[tensor] = tensor_val
+        scipy.io.savemat(self.results_save_dir+".mat", mdict=to_save_dict)
+            
+
+    # update learning rate (called once every epoch)
+    def update_learning_rate(self):
+        for scheduler in self.schedulers:
+            scheduler.step()
+        lr = self.optimizers[0].param_groups[0]['lr']
+        print('learning rate = %.7f' % lr)
+
+    # return visualization images. train.py will display these images, and save the images to a html
+    def get_current_visuals(self):
+        visual_ret = OrderedDict()
+        for name in self.visual_names:
+            if isinstance(name, str):
+                visual_ret[name] = getattr(self, name)
+        return visual_ret
+
+    # return traning losses/errors. train.py will print out these errors as debugging information
+    def get_current_losses(self):
+        errors_ret = OrderedDict()
+        for name in self.loss_names:
+            if isinstance(name, str):
+                # float(...) works for both scalar tensor and float number
+                errors_ret[name] = float(getattr(self, 'loss_' + name))
+        return errors_ret
+
+    # save models to the disk
+    def save_networks(self, epoch):
+        name = self.model_name
+        save_filename = '%s_net_%s.pth' % (epoch, name)
+        if not os.path.exists(self.save_dir):
+            os.makedirs(self.save_dir)
+        save_path = os.path.join(self.save_dir, save_filename)
+        net = self.net
+        print(save_path)
+        print('net'+name)
+        if len(self.gpu_ids) > 0 and torch.cuda.is_available() and hasattr(net, 'module'):
+            torch.save(net.module.cpu().state_dict(), save_path)
+        else:
+            torch.save(net.cpu().state_dict(), save_path)
+        net.to(self.device)
+
+        # save optimizer, scheduler, and RNG state
+        if self.isTrain:
+            optim_filename = '%s_optim_%s.pth' % (epoch, self.model_name)
+            optim_path = os.path.join(self.save_dir, optim_filename)
+            save_dict = {
+                'optimizer': self.optimizer.state_dict(),
+                'scheduler': self.schedulers[0].state_dict() if self.schedulers else None,
+                'rng_torch': torch.random.get_rng_state(),
+                'rng_numpy': np.random.get_state(),
+            }
+            if torch.cuda.is_available():
+                save_dict['rng_cuda'] = torch.cuda.get_rng_state_all()
+            torch.save(save_dict, optim_path)
+
+    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
+        key = keys[i]
+        if i + 1 == len(keys):  # at the end, pointing to a parameter/buffer
+            if module.__class__.__name__.startswith('InstanceNorm') and (key == 'running_mean' or key == 'running_var'):
+                if getattr(module, key) is None:
+                    state_dict.pop('.'.join(keys))
+            if module.__class__.__name__.startswith('InstanceNorm') and (key == 'num_batches_tracked'):
+                state_dict.pop('.'.join(keys))
+        else:
+            self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
+
+    # load models from the disk
+    def load_networks(self, epoch, load_dir=""):
+        """
+        epoch (int/str): epoch index / "best" / "latest"
+        """
+        assert isinstance(epoch,int) or epoch=="best" or epoch=="latest"
+        load_filename = f'{epoch}_net_{self.model_name}.pth'
+
+        if load_dir:
+            # use given load dir
+            load_path = os.path.join(load_dir, self.model_name, load_filename) 
+        else:
+            # use default load dir
+            load_path = os.path.join(self.load_dir, load_filename)
+
+        net = self.net
+        if isinstance(net, torch.nn.DataParallel):
+            net = net.module
+        print('loading the model from %s' % load_path)
+        # if you are using PyTorch newer than 0.4 (e.g., built from
+        # GitHub source), you can remove str() on self.device
+        state_dict = torch.load(load_path, weights_only=False)
+        if hasattr(state_dict, '_metadata'):
+            del state_dict._metadata
+
+        # patch InstanceNorm checkpoints prior to 0.4
+        print(state_dict.keys())
+        for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
+            self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
+        net.load_state_dict(state_dict)
+        net = net.to(self.device)
+
+    # load optimizer and scheduler state from disk
+    def load_training_state(self, epoch):
+        optim_filename = f'{epoch}_optim_{self.model_name}.pth'
+        load_path = os.path.join(self.load_dir, optim_filename)
+        if os.path.exists(load_path):
+            print(f'loading optimizer/scheduler from {load_path}')
+            state = torch.load(load_path, weights_only=False)
+            self.optimizer.load_state_dict(state['optimizer'])
+            if state.get('scheduler') and self.schedulers:
+                self.schedulers[0].load_state_dict(state['scheduler'])
+            # Restore RNG states for session continuity
+            if 'rng_torch' in state:
+                torch.random.set_rng_state(state['rng_torch'])
+            if 'rng_numpy' in state:
+                np.random.set_state(state['rng_numpy'])
+            if 'rng_cuda' in state and torch.cuda.is_available():
+                torch.cuda.set_rng_state_all(state['rng_cuda'])
+        else:
+            print(f'No optimizer state found at {load_path}, starting fresh')
+
+    # print network information
+    def print_networks(self, verbose):
+        print('---------- Networks initialized -------------')
+        net = self.net
+        name = self.model_name
+        num_params = 0
+        for param in net.parameters():
+            num_params += param.numel()
+        if verbose:
+            print(net)
+        print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
+        print('-----------------------------------------------')
+
+    # set requies_grad=Fasle to avoid computation
+    def set_requires_grad(self, nets, requires_grad=False):
+        if not isinstance(nets, list):
+            nets = [nets]
+        for net in nets:
+            if net is not None:
+                for param in net.parameters():
+                    param.requires_grad = requires_grad
+
+    # Set the input and target data for the network to train on/evaluate against
+    def set_data(self, input, target, convert=False, shuffle_channel=False):
+        shape_in = input.shape
+        if shuffle_channel:
+            self.input = input[:,torch.randperm(shape_in[1]),:,:]
+        else:
+            self.input = input
+        self.target = target
+        if convert:
+            self.input, self.target = Variable(self.input), Variable(self.target)
+
+        self.input = self.input.to(self.device)
+        self.target = self.target.to(self.device)
+
+    # Define the forward pass to compute loss
+    def forward(self):
+        self.output = self.net(self.input)
+        if self.opt.loss_type != "NoLoss":
+            self.loss = self.loss_weight*self.loss_criterion(self.output, self.target)
+
+            if self.opt.loss_type == "L1_sumL2" or self.opt.loss_type == "L2_sumL2":
+                self.loss += self.reg_loss_weight*torch.norm(self.output).div(self.output.numel())
+                self.reg_loss = self.reg_loss_weight*torch.norm(self.output).div(self.output.numel())
+
+            if self.opt.loss_type == "L2_sumL1":
+                self.loss += self.reg_loss_weight*torch.norm(self.output,p=1).div(self.output.numel())
+                self.reg_loss = self.reg_loss_weight*torch.norm(self.output,p=1).div(self.output.numel())
+
+            if self.opt.loss_type == "L1_sumL2_cross" or self.opt.loss_type == "L2_sumL2_cross":
+                self.loss += self.reg_loss_weight*torch.norm(self.output).div(self.output.numel())
+                self.loss += self.cross_loss_weight*self.cross_loss_criterion(self.output.flatten(start_dim=1),self.target.flatten(start_dim=1))
+            
+            if self.opt.loss_type == "L1_offset_loss" or self.opt.loss_type == "L2_offset_loss":
+                self.loss += self.reg_loss_weight*torch.norm(self.output).div(self.output.numel())
+
+    def backward(self):
+        self.loss.backward(retain_graph=True)
+    
+    def optimize_parameters(self):
+        self.forward()
+        self.backward()  
+        self.optimizer.step()
+        self.optimizer.zero_grad()

params.py

@@ -0,0 +1 @@
+from params_storage.params_fig10b_single_gpu import *

params_storage/params_fig10b_single_gpu.py

@@ -0,0 +1,178 @@
+# params_fig10b_single_gpu.py
+# ===========================
+# Reproduces Fig 10b from the paper:
+#   Complex environment - Jacobs Hall, July 28 + July 28_2, 4 APs
+#   Target: median ~64 cm, 90th ~160 cm
+#
+# Adapted for single GPU (RTX 3060 6GB) with batch_size=8.
+# To use: copy this file to ../params.py
+#   cp params_storage/params_fig10b_single_gpu.py params.py
+# Then run: python train_and_test.py
+#
+# Required data files in ./data/:
+#   dataset_jacobs_July28.mat
+#   dataset_fov_train_jacobs_July28_2.mat
+#   dataset_non_fov_train_jacobs_July28_2.mat
+#   dataset_fov_test_jacobs_July28_2.mat
+#   dataset_non_fov_test_jacobs_July28_2.mat
+# (Generate these with: python prepare_data.py)
+
+from easydict import EasyDict as edict
+import time
+from os.path import join
+
+opt_exp = edict()
+
+# ---------- Global Experiment params ----------
+opt_exp.isTrain = True
+opt_exp.continue_train = False
+opt_exp.starting_epoch_count = 0
+opt_exp.n_epochs = 50
+opt_exp.gpu_ids = ['0']           # Single GPU
+opt_exp.data = "rw_to_rw"        # Jacobs Hall complex env
+opt_exp.n_decoders = 2            # Encoder + location decoder + consistency decoder
+
+opt_exp.batch_size = 8            # Reduced from 32 for 6 GB VRAM
+opt_exp.ds_step_trn = 1
+opt_exp.ds_step_tst = 1
+opt_exp.weight_decay = 1e-5
+opt_exp.confidence = False
+
+# ------ Experiment save paths ----------
+opt_exp.save_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())  # No colons - Windows safe
+opt_exp.checkpoints_dir = join('./runs', opt_exp.save_name)
+opt_exp.results_dir = opt_exp.checkpoints_dir
+opt_exp.log_dir = opt_exp.checkpoints_dir
+opt_exp.load_dir = opt_exp.checkpoints_dir
+
+# ---------- Encoder params ----------
+opt_encoder = edict()
+opt_encoder.parent_exp = opt_exp
+opt_encoder.batch_size = opt_encoder.parent_exp.batch_size
+opt_encoder.ngf = 64
+opt_encoder.base_model = 'resnet_encoder'
+opt_encoder.net = 'G'
+opt_encoder.resnet_blocks = 6
+opt_encoder.no_dropout = False
+opt_encoder.init_type = 'xavier'
+opt_encoder.init_gain = 0.02
+opt_encoder.norm = 'instance'
+opt_encoder.beta1 = 0.5
+opt_encoder.lr = 0.00001
+opt_encoder.lr_policy = 'step'
+opt_encoder.lr_decay_iters = 50
+opt_encoder.lambda_L = 1
+opt_encoder.lambda_cross = 1e-5
+opt_encoder.lambda_reg = 5e-4
+opt_encoder.weight_decay = opt_encoder.parent_exp.weight_decay
+
+opt_encoder.input_nc = 4         # 4 APs in Jacobs Hall
+opt_encoder.output_nc = 1
+opt_encoder.save_latest_freq = 5000
+opt_encoder.save_epoch_freq = 1
+opt_encoder.n_epochs = opt_encoder.parent_exp.n_epochs
+opt_encoder.isTrain = True
+opt_encoder.continue_train = False
+opt_encoder.starting_epoch_count = opt_encoder.parent_exp.starting_epoch_count
+opt_encoder.name = 'encoder'
+opt_encoder.loss_type = "NoLoss"
+opt_encoder.niter = 20
+opt_encoder.niter_decay = 100
+
+opt_encoder.gpu_ids = opt_encoder.parent_exp.gpu_ids
+opt_encoder.num_threads = 4
+opt_encoder.checkpoints_load_dir = opt_encoder.parent_exp.load_dir
+opt_encoder.checkpoints_save_dir = opt_encoder.parent_exp.checkpoints_dir
+opt_encoder.results_dir = opt_encoder.parent_exp.results_dir
+opt_encoder.log_dir = opt_encoder.parent_exp.log_dir
+opt_encoder.max_dataset_size = float("inf")
+opt_encoder.verbose = False
+opt_encoder.suffix = ''
+
+# ---------- Location Decoder params ----------
+opt_decoder = edict()
+opt_decoder.parent_exp = opt_exp
+opt_decoder.batch_size = opt_decoder.parent_exp.batch_size
+opt_decoder.ngf = 64
+opt_decoder.base_model = 'resnet_decoder'
+opt_decoder.net = 'G'
+opt_decoder.resnet_blocks = 9
+opt_decoder.encoder_res_blocks = 6
+opt_decoder.no_dropout = False
+opt_decoder.init_type = 'xavier'
+opt_decoder.init_gain = 0.02
+opt_decoder.norm = 'instance'
+opt_decoder.beta1 = 0.5
+opt_decoder.lr = opt_encoder.lr
+opt_decoder.lr_policy = 'step'
+opt_decoder.lr_decay_iters = 20
+opt_decoder.lambda_L = 1
+opt_decoder.lambda_cross = 1e-5
+opt_decoder.lambda_reg = 5e-4
+opt_decoder.weight_decay = opt_decoder.parent_exp.weight_decay
+
+opt_decoder.input_nc = 4         # 4 APs
+opt_decoder.output_nc = 1        # 1-channel location heatmap
+opt_decoder.save_latest_freq = 5000
+opt_decoder.save_epoch_freq = 1
+opt_decoder.n_epochs = opt_decoder.parent_exp.n_epochs
+opt_decoder.isTrain = True
+opt_decoder.continue_train = False
+opt_decoder.starting_epoch_count = opt_decoder.parent_exp.starting_epoch_count
+opt_decoder.name = 'decoder'
+opt_decoder.loss_type = "L2_sumL1"
+opt_decoder.niter = 20
+opt_decoder.niter_decay = 100
+
+opt_decoder.gpu_ids = opt_decoder.parent_exp.gpu_ids
+opt_decoder.num_threads = 4
+opt_decoder.checkpoints_load_dir = opt_decoder.parent_exp.load_dir
+opt_decoder.checkpoints_save_dir = opt_decoder.parent_exp.checkpoints_dir
+opt_decoder.results_dir = opt_decoder.parent_exp.results_dir
+opt_decoder.log_dir = opt_decoder.parent_exp.log_dir
+opt_decoder.verbose = False
+opt_decoder.suffix = ''
+
+# ---------- Consistency (Offset) Decoder params ----------
+opt_offset_decoder = edict()
+opt_offset_decoder.parent_exp = opt_exp
+opt_offset_decoder.batch_size = opt_offset_decoder.parent_exp.batch_size
+opt_offset_decoder.ngf = 64
+opt_offset_decoder.base_model = 'resnet_decoder'
+opt_offset_decoder.net = 'G'
+opt_offset_decoder.resnet_blocks = 12
+opt_offset_decoder.encoder_res_blocks = 6
+opt_offset_decoder.no_dropout = False
+opt_offset_decoder.init_type = 'xavier'
+opt_offset_decoder.init_gain = 0.02
+opt_offset_decoder.norm = 'instance'
+opt_offset_decoder.beta1 = 0.5
+opt_offset_decoder.lr = opt_encoder.lr
+opt_offset_decoder.lr_policy = 'step'
+opt_offset_decoder.lr_decay_iters = 50
+opt_offset_decoder.lambda_L = 1
+opt_offset_decoder.lambda_cross = 0
+opt_offset_decoder.lambda_reg = 0
+opt_offset_decoder.weight_decay = opt_offset_decoder.parent_exp.weight_decay
+
+opt_offset_decoder.input_nc = 4  # 4 APs
+opt_offset_decoder.output_nc = 4 # 4-channel consistency output
+opt_offset_decoder.save_latest_freq = 5000
+opt_offset_decoder.save_epoch_freq = 1
+opt_offset_decoder.n_epochs = opt_offset_decoder.parent_exp.n_epochs
+opt_offset_decoder.isTrain = True
+opt_offset_decoder.continue_train = False
+opt_offset_decoder.starting_epoch_count = opt_offset_decoder.parent_exp.starting_epoch_count
+opt_offset_decoder.name = 'offset_decoder'
+opt_offset_decoder.loss_type = "L2_offset_loss"
+opt_offset_decoder.niter = 20
+opt_offset_decoder.niter_decay = 100
+
+opt_offset_decoder.gpu_ids = opt_offset_decoder.parent_exp.gpu_ids
+opt_offset_decoder.num_threads = 4
+opt_offset_decoder.checkpoints_load_dir = opt_offset_decoder.parent_exp.load_dir
+opt_offset_decoder.checkpoints_save_dir = opt_offset_decoder.parent_exp.checkpoints_dir
+opt_offset_decoder.results_dir = opt_offset_decoder.parent_exp.results_dir
+opt_offset_decoder.log_dir = opt_offset_decoder.parent_exp.log_dir
+opt_offset_decoder.verbose = False
+opt_offset_decoder.suffix = ''

params_storage/params_mamba_single_gpu.py

@@ -0,0 +1,92 @@
+# params_mamba_single_gpu.py
+# ===========================
+# Mamba-DLoc: replaces ResNet bottleneck with Mamba SSM blocks.
+# Single decoder (location only, no consistency decoder).
+#
+# To use: copy this file to ../params.py
+#   cp params_storage/params_mamba_single_gpu.py params.py
+# Then run: python train_and_test.py
+#
+# Required data files in ./data/ (same as Fig 10b):
+#   dataset_jacobs_July28.mat          (must be the FIXED 7,635-sample version)
+#   dataset_fov_train_jacobs_July28_2.mat
+#   dataset_non_fov_train_jacobs_July28_2.mat
+#   dataset_fov_test_jacobs_July28_2.mat
+#   dataset_non_fov_test_jacobs_July28_2.mat
+
+from easydict import EasyDict as edict
+import time
+from os.path import join
+
+opt_exp = edict()
+
+# ---------- Global Experiment params ----------
+opt_exp.isTrain = True
+opt_exp.continue_train = False
+opt_exp.starting_epoch_count = 0
+opt_exp.n_epochs = 50
+opt_exp.gpu_ids = ['0']
+opt_exp.data = "rw_to_rw"           # Same Jacobs Hall dataset as Fig 10b
+opt_exp.n_decoders = 1               # Single decoder (location only)
+opt_exp.model_type = "mamba"         # Selects Mamba pipeline in train_and_test.py
+
+opt_exp.batch_size = 8
+opt_exp.ds_step_trn = 1
+opt_exp.ds_step_tst = 1
+opt_exp.weight_decay = 1e-5
+opt_exp.confidence = False
+
+# ------ Experiment save paths ----------
+opt_exp.save_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
+opt_exp.checkpoints_dir = join('./runs', opt_exp.save_name)
+opt_exp.results_dir = opt_exp.checkpoints_dir
+opt_exp.log_dir = opt_exp.checkpoints_dir
+opt_exp.load_dir = opt_exp.checkpoints_dir
+
+# ---------- Mamba model params (single model, acts as both encoder+decoder) ----------
+opt_mamba = edict()
+opt_mamba.parent_exp = opt_exp
+opt_mamba.batch_size = opt_exp.batch_size
+opt_mamba.ngf = 64
+opt_mamba.base_model = 'mamba_dloc'
+opt_mamba.net = 'G'
+opt_mamba.no_dropout = False
+opt_mamba.init_type = 'xavier'
+opt_mamba.init_gain = 0.02
+opt_mamba.norm = 'batch'
+opt_mamba.beta1 = 0.5
+opt_mamba.lr = 0.00001
+opt_mamba.lr_policy = 'step'
+opt_mamba.lr_decay_iters = 20
+opt_mamba.lambda_L = 1
+opt_mamba.lambda_cross = 0
+opt_mamba.lambda_reg = 5e-4
+opt_mamba.weight_decay = opt_exp.weight_decay
+
+opt_mamba.input_nc = 4
+opt_mamba.output_nc = 1
+opt_mamba.save_latest_freq = 5000
+opt_mamba.save_epoch_freq = 1
+opt_mamba.n_epochs = opt_exp.n_epochs
+opt_mamba.isTrain = True
+opt_mamba.continue_train = False
+opt_mamba.starting_epoch_count = opt_exp.starting_epoch_count
+opt_mamba.name = 'mamba_dloc'
+opt_mamba.loss_type = "L2_sumL1"
+opt_mamba.niter = 20
+opt_mamba.niter_decay = 100
+
+opt_mamba.gpu_ids = opt_exp.gpu_ids
+opt_mamba.num_threads = 4
+opt_mamba.checkpoints_load_dir = opt_exp.load_dir
+opt_mamba.checkpoints_save_dir = opt_exp.checkpoints_dir
+opt_mamba.results_dir = opt_exp.results_dir
+opt_mamba.log_dir = opt_exp.log_dir
+opt_mamba.max_dataset_size = float("inf")
+opt_mamba.verbose = False
+opt_mamba.suffix = ''
+
+# ---------- Aliases for train_and_test.py compatibility ----------
+opt_encoder = opt_mamba
+opt_decoder = opt_mamba
+opt_offset_decoder = opt_mamba

requirements.txt

@@ -0,0 +1,6 @@
+easydict>=1.9
+h5py>=3.1.0
+hdf5storage>=0.1.15
+numpy>=1.21.0
+scipy>=1.7.0
+scikit-learn

train_and_test.py

@@ -0,0 +1,226 @@
+#!/usr/bin/python
+'''
+Script for both training and evaluating the DLoc network
+Automatically imports the parameters from params.py.
+For further details onto which params file to load
+read the README in `params_storage` folder.
+'''
+import sys
+import torch
+import warnings
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore",category=FutureWarning)
+from utils import *
+from modelADT import ModelADT
+from Generators import *
+from data_loader import load_data, LazyMultiHDF5Dataset
+from joint_model import Enc_2Dec_Network
+from joint_model import Enc_Dec_Network
+from joint_model import Mamba_Network
+from params import *
+import trainer
+torch.manual_seed(0)
+np.random.seed(0)
+
+'''
+Defining the paths from where to Load Data.
+Assumes that the data is stored in a subfolder called data in the current data folder
+'''
+
+#####################################Final Simple Space Results################################################
+if "data" in opt_exp and opt_exp.data == "rw_to_rw_atk":
+    # Training and testing data loaded for the Final results For Env-1 (The smaller space) in the paper (Figure 10a)
+    trainpath = ['./data/dataset_non_fov_train_July18.mat',
+                './data/dataset_fov_train_July18.mat']
+    testpath = ['./data/dataset_non_fov_test_July18.mat',
+                './data/dataset_fov_test_July18.mat']
+    print('Real World to Real World experiments started')
+
+#####################################Final Complex Space Results################################################
+elif "data" in opt_exp and opt_exp.data == "rw_to_rw":
+    # Training and testing data loaded for the Final results For Env-2 (The larger space) in the paper (Figure 10b)
+    trainpath = ['./data/dataset_jacobs_July28.mat',
+                './data/dataset_non_fov_train_jacobs_July28_2.mat',
+                './data/dataset_fov_train_jacobs_July28_2.mat']
+    testpath = ['./data/dataset_fov_test_jacobs_July28_2.mat',
+                './data/dataset_non_fov_test_jacobs_July28_2.mat']
+    print('Real World to Real World experiments started')
+
+#########################################Generalization across Scenarios###########################################
+
+elif "data" in opt_exp and opt_exp.data == "rw_to_rw_env2":
+    # Training and testing data loaded for the Final results For Env-2
+    # for Generalization across scenarios (Table-1) train on 1/3/4 and test on 2
+    trainpath = ['./data/dataset_jacobs_July28.mat',
+                './data/dataset_non_fov_train_jacobs_July28_2.mat',
+                './data/dataset_fov_train_jacobs_July28_2.mat',
+                './data/dataset_train_jacobs_Aug16_3.mat',
+                './data/dataset_train_jacobs_Aug16_4_ref.mat']
+    testpath = ['./data/dataset_train_jacobs_Aug16_1.mat']
+    print('Real World to Real World experiments started')
+
+
+elif "data" in opt_exp and opt_exp.data == "rw_to_rw_env3":
+    # Training and testing data loaded for the Final results For Env-2
+    # for Generalization across scenarios (Table-1) train on 1/2/4 and test on 3
+    trainpath = ['./data/dataset_jacobs_July28.mat',
+                './data/dataset_non_fov_train_jacobs_July28_2.mat',
+                './data/dataset_fov_train_jacobs_July28_2.mat',
+                './data/dataset_train_jacobs_Aug16_1.mat',
+                './data/dataset_train_jacobs_Aug16_4_ref.mat']
+    testpath = ['./data/dataset_train_jacobs_Aug16_3.mat']
+    print('Real World to Real World experiments started')
+
+elif "data" in opt_exp and opt_exp.data == "rw_to_rw_env4":
+    # Training and testing data loaded for the Final results For Env-2
+    # for Generalization across scenarios (Table-1) train on 1/2/3 and test on 4
+    trainpath = ['./data/dataset_jacobs_July28.mat',
+                './data/dataset_non_fov_train_jacobs_July28_2.mat',
+                './data/dataset_fov_train_jacobs_July28_2.mat',
+                './data/dataset_train_jacobs_Aug16_1.mat',
+                './data/dataset_train_jacobs_Aug16_3.mat']
+    testpath = ['./data/dataset_train_jacobs_Aug16_4_ref.mat']
+    print('Real World to Real World experiments started')
+
+######################################Generalization Across Bandwidth##########################################
+
+elif "data" in opt_exp and opt_exp.data == "rw_to_rw_40":
+    # Training and testing data loaded for the Generalization results For Env-2 (The larger space) in the paper (Figure 13a) at 40MHz
+    trainpath = ['./data/dataset40_jacobs_July28.mat',
+                './data/dataset40_non_fov_train_jacobs_July28_2.mat',
+                './data/dataset40_fov_train_jacobs_July28_2.mat']
+    testpath = ['./data/dataset40_fov_test_jacobs_July28_2.mat',
+                './data/dataset40_non_fov_test_jacobs_July28_2.mat']
+    print('Real World to Real World experiments started')
+
+elif "data" in opt_exp and opt_exp.data == "rw_to_rw_20":
+    # Training and testing data loaded for the Generalization results For Env-2 (The larger space) in the paper (Figure 13a) at 20MHz
+    trainpath = ['./data/dataset20_jacobs_July28.mat',
+                './data/dataset20_non_fov_train_jacobs_July28_2.mat',
+                './data/dataset20_fov_train_jacobs_July28_2.mat']
+    testpath = ['./data/dataset20_fov_test_jacobs_July28_2.mat',
+                './data/dataset20_non_fov_test_jacobs_July28_2.mat']
+    print('Real World to Real World experiments started')
+
+######################################Generalization Across Space##########################################
+
+elif "data" in opt_exp and opt_exp.data == "data_segment":
+    # Training and testing data loaded for the Final results For Env-2 
+    # for Disjoint Training and Testing(The larger space) in the paper (Figure 13b)
+    trainpath = ['./data/dataset_train_jacobs_July28.mat',
+                './data/dataset_train_jacobs_July28_2.mat']
+    testpath = ['./data/dataset_test_jacobs_July28.mat',
+                './data/dataset_test_jacobs_July28_2.mat']
+    print('non-FOV to non-FOV experiments started')
+
+######################################################################################################################
+'''
+Loading Training and Evaluation Data via lazy disk-backed Dataset.
+Samples are read on demand from HDF5 files - uses ~0 RAM.
+'''
+print('Indexing training files (lazy, no data loaded into RAM)...')
+train_data = LazyMultiHDF5Dataset(trainpath)
+train_loader = torch.utils.data.DataLoader(
+    train_data, batch_size=opt_exp.batch_size, shuffle=True, num_workers=0)
+print(f'# training samples = {len(train_data)}')
+print(f'# training mini batch = {len(train_loader)}')
+
+print('Indexing testing files...')
+test_data = LazyMultiHDF5Dataset(testpath)
+test_loader = torch.utils.data.DataLoader(
+    test_data, batch_size=opt_exp.batch_size, shuffle=False, num_workers=0)
+print(f'# testing samples = {len(test_data)}')
+print(f'# testing mini batch = {len(test_loader)}')
+print('Data Loaded (lazy)')
+
+'''
+Initiate the Network and build the graph
+'''
+
+model_type = getattr(opt_exp, 'model_type', 'dloc')
+
+if model_type == 'mamba':
+    # Mamba: single end-to-end model wrapped in one ModelADT
+    from params import opt_mamba
+    mamba_model = ModelADT()
+    mamba_model.initialize(opt_mamba)
+    mamba_model.setup(opt_mamba)
+
+    print('Making the Mamba joint model')
+    joint_model = Mamba_Network()
+    joint_model.initialize(opt_exp, mamba_model, gpu_ids=opt_exp.gpu_ids)
+
+else:
+    # Original DLoc pipeline
+    # init encoder
+    enc_model = ModelADT()
+    enc_model.initialize(opt_encoder)
+    enc_model.setup(opt_encoder)
+
+    # init decoder1
+    dec_model = ModelADT()
+    dec_model.initialize(opt_decoder)
+    dec_model.setup(opt_decoder)
+
+    if opt_exp.n_decoders == 2:
+        # init decoder2
+        offset_dec_model = ModelADT()
+        offset_dec_model.initialize(opt_offset_decoder)
+        offset_dec_model.setup(opt_offset_decoder)
+
+        # join all models
+        print('Making the joint_model')
+        joint_model = Enc_2Dec_Network()
+        joint_model.initialize(opt_exp, enc_model, dec_model, offset_dec_model, gpu_ids=opt_exp.gpu_ids)
+
+    elif opt_exp.n_decoders == 1:
+        # join all models
+        print('Making the joint_model')
+        joint_model = Enc_Dec_Network()
+        joint_model.initialize(opt_exp, enc_model, dec_model, gpu_ids=opt_exp.gpu_ids)
+
+    else:
+        print('Incorrect number of Decoders specified in the parameters')
+        sys.exit(-1)
+
+    if getattr(opt_exp, 'isFrozen', False):
+        enc_model.load_networks(opt_encoder.starting_epoch_count)
+        dec_model.load_networks(opt_decoder.starting_epoch_count)
+        if opt_exp.n_decoders == 2:
+            offset_dec_model.load_networks(opt_offset_decoder.starting_epoch_count)
+
+# train the model
+'''
+Trainig the network
+'''
+trainer.train(joint_model, train_loader, test_loader)
+
+'''
+Model Evaluation at the best epoch
+'''
+
+epoch = "best"  # int/"best"/"last"
+eval_name = opt_exp.checkpoints_dir
+
+if model_type == 'mamba':
+    # Mamba: reload single model
+    mamba_model.load_networks(epoch, load_dir=eval_name)
+    joint_model.initialize(opt_exp, mamba_model, gpu_ids=opt_exp.gpu_ids)
+else:
+    # DLoc: reload encoder + decoder(s)
+    enc_model.load_networks(epoch, load_dir=eval_name)
+    dec_model.load_networks(epoch, load_dir=eval_name)
+    if opt_exp.n_decoders == 2:
+        offset_dec_model.load_networks(epoch, load_dir=eval_name)
+        joint_model.initialize(opt_exp, enc_model, dec_model, offset_dec_model, gpu_ids=opt_exp.gpu_ids)
+    elif opt_exp.n_decoders == 1:
+        joint_model.initialize(opt_exp, enc_model, dec_model, gpu_ids=opt_exp.gpu_ids)
+
+# pass data through model
+total_loss, median_error = trainer.test(joint_model,
+    test_loader,
+    save_output=True,
+    save_dir=eval_name,
+    save_name=f"decoder_test_result_epoch_{epoch}",
+    log=False)
+print(f"total_loss: {total_loss}, median_error: {median_error}")

trainer.py

@@ -0,0 +1,182 @@
+'''
+Scripts for the training and testing functions
+train() function is called for training the network
+test() function is called to evaluate the network
+Both the function logs and saves the results in the files 
+as mentioned in the params.py file
+'''
+import warnings
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore",category=FutureWarning)
+import numpy as np  
+import os
+import time
+import hdf5storage
+from utils import *
+from Generators import *
+from params import *
+
+def train(model, train_loader, test_loader):
+    """Traning pipeline
+
+    Args:
+        model (torch.module): pytorch model
+        train_loader (torch.dataloader): dataloader
+        test_loader (torch.dataloader): dataloader
+    """
+    # set data index
+    offset_output_index=0
+    input_index=1
+    output_index=2
+    
+    # initialization
+    total_steps = 0
+    print('Training called')
+    stopping_count = 0
+
+    for epoch in range(model.opt.starting_epoch_count+1, model.opt.n_epochs+1): # opt.niter + opt.niter_decay + 1):
+        epoch_start_time = time.time()
+        epoch_loss = 0
+        epoch_offset_loss = 0
+        error =[]
+
+        for i, data in enumerate(train_loader):
+            total_steps += model.opt.batch_size
+            if opt_exp.n_decoders == 2:
+                model.set_input(data[input_index], data[output_index], data[offset_output_index], shuffle_channel=False)
+            elif opt_exp.n_decoders == 1:
+                model.set_input(data[input_index], data[output_index], shuffle_channel=False)
+            model.optimize_parameters()
+            dec_outputs = model.decoder.output
+            # print(f"dec_outputs size is : {dec_outputs.shape}")
+            error.extend(localization_error(dec_outputs.data.cpu().numpy(),data[output_index].cpu().numpy(),scale=0.1))
+
+            write_log([str(model.decoder.loss.item())], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='loss')
+            if opt_exp.n_decoders == 2:
+                write_log([str(model.offset_decoder.loss.item())], model.offset_decoder.model_name, log_dir=model.offset_decoder.opt.log_dir, log_type='offset_loss')
+            if total_steps % model.decoder.opt.save_latest_freq == 0:
+                print('saving the latest model (epoch %d, total_steps %d)' %
+                      (epoch, total_steps))
+                model.save_networks('latest')
+
+            epoch_loss += model.decoder.loss.item()
+            if opt_exp.n_decoders == 2:
+                epoch_offset_loss += model.offset_decoder.loss.item()
+
+        median_error_tr = np.median(error)
+        error_90th_tr = np.percentile(error,90)
+        error_99th_tr = np.percentile(error,99)
+        nighty_percentile_error_tr = np.percentile(error,90)
+        epoch_loss /= i
+        if opt_exp.n_decoders == 2:
+            epoch_offset_loss /= i
+        write_log([str(epoch_loss)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='epoch_decoder_loss')
+        if opt_exp.n_decoders == 2:
+            write_log([str(epoch_offset_loss)], model.offset_decoder.model_name, log_dir=model.offset_decoder.opt.log_dir, log_type='epoch_offset_decoder_loss')
+        write_log([str(median_error_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_median_error')
+        write_log([str(error_90th_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_90th_error')
+        write_log([str(error_99th_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_99th_error')
+        write_log([str(nighty_percentile_error_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_90_error')
+        if (epoch == model.opt.starting_epoch_count + 1):
+            min_eval_loss, median_error = test(model, test_loader, save_output=False)
+        else:
+            new_eval_loss, new_med_error = test(model, test_loader, save_output=False)
+            if (median_error>=new_med_error):
+                stopping_count = stopping_count+1
+                median_error = new_med_error
+
+        # generated_outputs = temp_generator_outputs
+        if epoch % model.encoder.opt.save_epoch_freq == 0:
+            print('saving the model at the end of epoch %d, iters %d' %(epoch, total_steps))
+            model.save_networks('latest')
+            model.save_networks(epoch)
+            if (stopping_count==2):
+                print('Saving best model at %d epoch' %(epoch))
+                model.save_networks('best')
+                stopping_count=0
+
+        print('End of epoch %d / %d \t Time Taken: %d sec' %
+              (epoch, model.decoder.opt.niter + model.decoder.opt.niter_decay, time.time() - epoch_start_time))
+        model.decoder.update_learning_rate()
+        model.encoder.update_learning_rate()
+        if opt_exp.n_decoders == 2:
+            model.offset_decoder.update_learning_rate()
+
+
+def test(model, test_loader, save_output=True, save_name="decoder_test_result", save_dir="", log=True):
+    """Test and evaluation pipeline
+
+    Args:
+        model (torch.module): pytorch model
+        test_loader (torch.dataloader): dataloader
+        save_output (bool, optional): whether to save output to mat file. Defaults to True.
+        save_name (str, optional): name of the mat file. Defaults to "decoder_test_result".
+        save_dir (str, optional): directory where output mat file is saved. Defaults to "".
+        log (bool, optional): whether to log output. Defaults to True.
+
+    Returns:
+        tuple: (total_loss -> float, median_error -> float)
+    """
+    print('Evaluation Called')
+    model.eval()
+
+    # set data index
+    offset_output_index=0
+    input_index=1
+    output_index=2
+
+    # create containers
+    generated_outputs = []
+    offset_outputs = []
+    total_loss = 0
+    total_offset_loss = 0
+    error =[]
+    for i, data in enumerate(test_loader):
+        if opt_exp.n_decoders == 2:
+            model.set_input(data[input_index], data[output_index], data[offset_output_index], shuffle_channel=False)
+        elif opt_exp.n_decoders == 1:
+            model.set_input(data[input_index], data[output_index], shuffle_channel=False)
+        model.test()
+
+        # get model outputs
+        gen_outputs = model.decoder.output  # gen_outputs.size = (N,1,H,W)
+        if opt_exp.n_decoders == 2:
+            off_outputs = model.offset_decoder.output # off_outputs.size = (N,n_ap,H,W)
+
+        generated_outputs.extend(gen_outputs.data.cpu().numpy())
+        if opt_exp.n_decoders == 2:
+            offset_outputs.extend(off_outputs.data.cpu().numpy())
+        error.extend(localization_error(gen_outputs.data.cpu().numpy(),data[output_index].cpu().numpy(),scale=0.1))
+        total_loss += model.decoder.loss.item()
+        if opt_exp.n_decoders == 2:
+            total_offset_loss += model.offset_decoder.loss.item()
+    total_loss /= i
+    if opt_exp.n_decoders == 2:
+        total_offset_loss /= i
+    median_error = np.median(error)
+    nighty_percentile_error = np.percentile(error,90)
+    error_99th = np.percentile(error,99)
+
+    if log:
+        write_log([str(median_error)], model.decoder.model_name, log_dir=model.opt.log_dir, log_type='test_median_error')
+        write_log([str(nighty_percentile_error)], model.decoder.model_name, log_dir=model.opt.log_dir, log_type='test_90_error')
+        write_log([str(error_99th)], model.decoder.model_name, log_dir=model.opt.log_dir, log_type='test_99_error')
+        write_log([str(total_loss)], model.decoder.model_name, log_dir=model.opt.log_dir, log_type='test_loss')
+        if opt_exp.n_decoders == 2:
+            write_log([str(total_offset_loss)], model.offset_decoder.model_name, log_dir=model.opt.log_dir, log_type='test_offset_loss')
+
+    if save_output:
+        if not save_dir:
+            save_dir = model.decoder.results_save_dir # default save directory
+
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+        
+        save_path = f"{save_dir}/{save_name}.mat"
+        hdf5storage.savemat(save_path,
+            mdict={"outputs":generated_outputs,"wo_outputs":offset_outputs, "error": error}, 
+            appendmat=True, 
+            format='7.3',
+            truncate_existing=True)
+        print(f"result saved in {save_path}")
+    return total_loss, median_error

utils.py

@@ -0,0 +1,149 @@
+#!/usr/bin/python
+'''
+Contains the utilities used for
+loading, initating and running up the networks
+for all training, validation and evaluation.
+'''
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from torch.optim import lr_scheduler
+import numpy as np
+import os
+from Generators import *
+from mamba_model import MambaDLocNet
+
+def write_log(log_values, model_name, log_dir="", log_type='loss', type_write='a'):
+    if not os.path.exists(log_dir):
+            os.makedirs(log_dir)
+    with open(log_dir+"/"+model_name+"_"+log_type+".txt", type_write) as f:
+        f.write(','.join(log_values)+"\n")
+
+def get_model_funct(model_name):
+    if model_name == "G":
+        return define_G
+
+def define_G(opt, gpu_ids):
+    net = None
+    input_nc    = opt.input_nc
+    output_nc   = opt.output_nc
+    ngf         = opt.ngf
+    net_type    = opt.base_model
+    norm        = opt.norm
+    use_dropout = opt.no_dropout
+    init_type   = opt.init_type
+    init_gain   = opt.init_gain
+
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if net_type == 'resnet_encoder':
+        n_blocks    = opt.resnet_blocks
+        net = ResnetEncoder(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=n_blocks)
+    elif net_type == 'resnet_decoder':
+        n_blocks    = opt.resnet_blocks
+        net = ResnetDecoder(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=n_blocks, encoder_blocks=opt.encoder_res_blocks)
+    elif net_type == 'mamba_dloc':
+        net = MambaDLocNet(input_nc=input_nc)
+    else:
+        raise NotImplementedError('Generator model name [%s] is not recognized' % net_type)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+def get_scheduler(optimizer, opt):
+    if opt.starting_epoch_count=='best' and opt.lr_policy == 'lambda':
+        def lambda_rule(epoch):
+            print("lambda update %s, %s, %s", (epoch, 0))
+            lr_l = 1.0 - max(0, epoch + 1 - opt.niter) / float(opt.niter_decay + 1)
+            return lr_l
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif opt.starting_epoch_count!='best' and opt.lr_policy == 'lambda':
+        def lambda_rule(epoch):
+            print("lambda update %s, %s, %s", (epoch, opt.starting_epoch_count))
+            lr_l = 1.0 - max(0, epoch + 1 + opt.starting_epoch_count - opt.niter) / float(opt.niter_decay + 1)
+            return lr_l
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif opt.starting_epoch_count!='best' and opt.lr_policy == 'step':
+        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.9)
+    elif opt.starting_epoch_count!='best' and opt.lr_policy == 'plateau':
+        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
+    elif opt.starting_epoch_count!='best' and opt.lr_policy == 'cosine':
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.niter, eta_min=0)
+    else:
+        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
+    return scheduler
+
+
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+    elif norm_type == 'none':
+        norm_layer = None
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+
+def init_weights(net, init_type='normal', gain=1):
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0, gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)
+
+
+def init_net(net, init_type='normal', init_gain=1, gpu_ids=[]):
+    if len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        device_ = torch.device('cuda:{}'.format(gpu_ids[0]))
+#         net.to(d)
+        gpu_ids_int = list(map(int,gpu_ids))
+        net = torch.nn.DataParallel(net, gpu_ids_int)
+        net.to(device_)
+    init_weights(net, init_type, gain=init_gain)
+    return net
+
+def localization_error(output_predictions,input_labels,scale=0.1):
+    """
+    output_predictions: (N,1,H,W), model prediction
+    input_labels: (N,1,H,W), ground truth target
+    """
+    image_size = output_predictions.shape
+    error = np.zeros(image_size[0])
+
+    for i in range(image_size[0]):
+        label_temp = input_labels[i,:,:,:].squeeze() # ground truth label
+        pred_temp = output_predictions[i,:,:,:].squeeze() # model prediction
+        label_index = np.asarray(np.unravel_index(np.argmax(label_temp), label_temp.shape))
+        prediction_index = np.asarray(np.unravel_index(np.argmax(pred_temp),pred_temp.shape))
+        error[i] = np.sqrt( np.sum( np.power(np.multiply( label_index-prediction_index, scale ), 2)) )
+
+    return error
+
+class Flatten(nn.Module):
+    def __init__(self):
+        super(Flatten, self).__init__()
+
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+def RGB2Gray(img):
+    return 0.2125*img[:,:,0] + 0.7154*img[:,:,1] + 0.0721*img[:,:,2]