model.classify package

Submodules

model.classify.dataloader module

class model.classify.dataloader.EMRIDatasetTorch(*args: Any, **kwargs: Any)

Bases: Dataset

class model.classify.dataloader.NpzDatasetTorch(*args: Any, **kwargs: Any)

Bases: Dataset

class model.classify.dataloader.TinyEMRIDataset(DIR, fn)

Bases: object

save(DIR, fn)
model.classify.dataloader.load_waveform(data=None, DIR='.', data_fn='waveform_dataset.hdf5')

load waveform dataset from hdf5 file.

Parameters:

  • DIR (str, optional) – Specified directory to the waveform dataset, by default ‘.’

  • data_fn (str, optional) – Specified file name to the waveform dataset, by default ‘waveform_dataset.hdf5’

model.classify.dataloader.save_waveform(data=None, DIR='.', data_fn='waveform_dataset.hdf5')

Save waveform dataset to hdf5 file.

Parameters:

  • data (list, optional) – Dataset to save, [waveform_dataset, waveform_params]

  • DIR (str, optional) – Specified directory to save waveform dataset, by default ‘.’

  • data_fn (str, optional) – Specified file name to save waveform dataset, by default ‘waveform_dataset.hdf5’

model.classify.tcn module

class model.classify.tcn.Chomp1d(*args: Any, **kwargs: Any)

Bases: Module

forward(x)

Forward pass

Parameters:

x – Input data

Returns:

Cropped data

class model.classify.tcn.TCN(*args: Any, **kwargs: Any)

Bases: Module

clear_activation_maps()

Clear the activation maps

forward(inputs)

Forward pass

Parameters:

inputs – Input data dimension (N, C_in, L_in)

get_activation_maps()

Get the activation maps

class model.classify.tcn.TemporalBlock(*args: Any, **kwargs: Any)

Bases: Module

forward(x)

Forward pass

Parameters:

x – Input data

init_weights()

Initialize weights

class model.classify.tcn.TemporalConvNet(*args: Any, **kwargs: Any)

Bases: Module

clear_activation_maps()

Clear the activation maps

forward(x)

Forward pass

Parameters:

x – Input data

get_activation_maps()

Get the activation maps

hook_fn(module, input, output)

Hook function to store the activation maps

Parameters:

  • module – Module

  • input – Input

  • output – Output

model.classify.trainer module

class model.classify.trainer.Trainer(model, train_loader, test_loader, optimizer_type, optimizer_kwargs, scheduler_type, scheduler_kwargs, loss_fn, device, result_dir, result_fn='inf_result.npy', checkpoint_dir=None, use_wandb=False)

Bases: object

accuracy(output, target)

Compute the accuracy of the model.

Parameters:

  • output (torch.Tensor) – The model output.

  • target (torch.Tensor) – The target labels.

Returns:

The accuracy of the model.

Return type:

float

calculate_accuracy(logits, true_labels, threshold=0.5)

Calculate the classification accuracy for binary classification.

Parameters:

  • logits – The raw output scores from the model (before sigmoid). Expected shape is (batch_size).

  • true_labels – The ground truth labels. Expected shape is (batch_size).

  • threshold – The threshold to classify samples as positive or negative.

Returns:

The accuracy as a floating point number between 0 and 1.

Return type:

float

count_parameters(model)

Count the number of parameters in the model.

Parameters:

model – torch.nn.Module The model to count the parameters of.

evaluate(dataloader, save=False)

Evaluate the model on the test data.

Parameters:

  • dataloader – torch.utils.data.DataLoader The test data loader.

  • save – bool Whether to save the inference results.

Returns:

float The average test loss. avg_acc : float The average test accuracy.

Return type:

avg_loss

format_number_with_unit(num)

Format a number with a unit.

Parameters:

num – int The number to format.

Returns:

The formatted number string.

Return type:

str

get_gpu_tensors_sorted_by_size()

Get the GPU tensors sorted by size.

get_loss_fn(loss_choice, **kwargs)

Get the loss function.

Parameters:

  • loss_choice – str The loss function to use.

  • **kwargs – dict The loss function keyword arguments.

Returns:

The loss function.

Return type:

torch.nn.modules.loss._Loss

get_lr_scheduler(scheduler_choice, optimizer, **kwargs)

Get the learning rate scheduler.

Parameters:

  • scheduler_choice – str The scheduler to use.

  • optimizer – torch.optim.Optimizer The optimizer.

  • **kwargs – dict The scheduler keyword arguments.

Returns:

The learning rate scheduler.

Return type:

torch.optim.lr_scheduler._LRScheduler

get_optimizer(optimizer_choice, model, **kwargs)

Get the optimizer.

Parameters:

  • optimizer_choice – str The optimizer to use.

  • model – torch.nn.Module The model to optimize.

  • **kwargs – dict The optimizer keyword arguments.

Returns:

The optimizer.

Return type:

torch.optim.Optimizer

load_checkpoint()

Load the model from the checkpoint.

print_gpu_tensors(top_n=None)

Print the GPU tensors sorted by size.

Parameters:

top_n – int The number of tensors to display.

train(n_epochs)

Train the model for n_epochs.

Parameters:

n_epochs – int The number of epochs to train the model.

train_epoch(dataloader, epoch)

Train the model for one epoch.

Parameters:

  • dataloader – torch.utils.data.DataLoader The training data loader.

  • epoch – int The current epoch number.

Returns:

float The average training loss. avg_acc : float The average training accuracy.

Return type:

avg_loss

update_average(loss, avg_loss)

Update running average of the loss.

Args: loss : torch.tensor

detached loss, a single float value.

avg_lossfloat

current running average.

Returns: avg_loss : float

The average loss.

Module contents

Detect EMRI signal using DNN

model.classify.get_version() → str