<gh_stars>1-10
"""
Contains generic pytorch model operations such as loading,
forward and backward passes and tensor evaluations.
"""

import numpy as np
from torch import nn
import torch

#TODO add typing and sphinx docstrings

def set_seed(seed: int):
    """
    Sets random seed for numpy and both CPU/GPU pytorch routines.
    If cuda is not available, this is silently ignored.
    Note: numpy and pytorch do NOT use the same RNG.
    """
    np.random.seed(seed)
    torch.manual_seed(seed)

def load_model(model: nn.Module, model_args: dict, model_weights: str, device: torch.device):
    """
    Given model class, instantiation arguments and weight file, load the model
    """
    model = model(**model_args)
    state_dict = torch.load(model_weights, map_location=device)
    model.load_state_dict(state_dict["model"])
    return model

def forward(model: nn.Module, inputs: torch.Tensor, device: torch.device):
    """
    Performs a feed forward step of the
    given model and returns model outputs given a torch.Tensor input.
    """

    model.eval()
    model.to(device)

    with torch.no_grad():
        inputs = inputs.to(device)
        return model(inputs)

def backprop(model, loss_function, optimizer, batch, device):
    """
    Performs a single step of backpropagation
    and weight optimization
    """
    model.train()
    model.to(device)
    optimizer.zero_grad()

    inputs, targets = batch[0], batch[1]

    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()

def assert_not_nan(tensor):
    assert not torch.isnan(tensor).any()

def assert_isfinite(tensor):
    assert torch.isfinite(tensor).all()

def assert_upperbound(tensor, val):
    assert (tensor < val).all()

def assert_lowerbound(tensor, val):
    assert (tensor > val).all()

def assert_nonzero(tensor):
    assert (tensor != 0).all()
    
def set_module_trainable(module: nn.Module, mode: bool) -> None:
    """Freeze or unfreeze weights"""
    for param in module.parameters():
        param.requires_grad = mode
    
def visualise(x):
    """Slice and normalise a batch of images for visualisation"""
    assert len(x.shape) == 4
    if x.shape[1] == 2: # Sentinel 1
        return x[:, [0], :, :]
    elif x.shape[1] == 13: # Sentinel 2
        x = x[:, [3, 2, 1], :, :]
        x = (x - x.min()) / (x.max() - x.min() + 1e-5)
        return x**0.5 # gamma correction
    elif x.shape[1] > 3:
        return x[:,:3, :, :]
    else:
        return x
"""
Contains generic pytorch model operations such as loading,
forward and backward passes and tensor evaluations.
"""

import numpy as np
from torch import nn
import torch

#TODO add typing and sphinx docstrings

def set_seed(seed: int):
    """
    Sets random seed for numpy and both CPU/GPU pytorch routines.
    If cuda is not available, this is silently ignored.
    Note: numpy and pytorch do NOT use the same RNG.
    """
    np.random.seed(seed)
    torch.manual_seed(seed)

def load_model(model: nn.Module, model_args: dict, model_weights: str, device: torch.device):
    """
    Given model class, instantiation arguments and weight file, load the model
    """
    model = model(**model_args)
    state_dict = torch.load(model_weights, map_location=device)
    model.load_state_dict(state_dict["model"])
    return model

def forward(model: nn.Module, inputs: torch.Tensor, device: torch.device):
    """
    Performs a feed forward step of the
    given model and returns model outputs given a torch.Tensor input.
    """

    model.eval()
    model.to(device)

    with torch.no_grad():
        inputs = inputs.to(device)
        return model(inputs)

def backprop(model, loss_function, optimizer, batch, device):
    """
    Performs a single step of backpropagation
    and weight optimization
    """
    model.train()
    model.to(device)
    optimizer.zero_grad()

    inputs, targets = batch[0], batch[1]

    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()

def assert_not_nan(tensor):
    assert not torch.isnan(tensor).any()

def assert_isfinite(tensor):
    assert torch.isfinite(tensor).all()

def assert_upperbound(tensor, val):
    assert (tensor < val).all()

def assert_lowerbound(tensor, val):
    assert (tensor > val).all()

def assert_nonzero(tensor):
    assert (tensor != 0).all()
    
def set_module_trainable(module: nn.Module, mode: bool) -> None:
    """Freeze or unfreeze weights"""
    for param in module.parameters():
        param.requires_grad = mode
    
def visualise(x):
    """Slice and normalise a batch of images for visualisation"""
    assert len(x.shape) == 4
    if x.shape[1] == 2: # Sentinel 1
        return x[:, [0], :, :]
    elif x.shape[1] == 13: # Sentinel 2
        x = x[:, [3, 2, 1], :, :]
        x = (x - x.min()) / (x.max() - x.min() + 1e-5)
        return x**0.5 # gamma correction
    elif x.shape[1] > 3:
        return x[:,:3, :, :]
    else:
        return x
Contains generic pytorch model operations such as loading, forward and backward passes and tensor evaluations. #TODO add typing and sphinx docstrings Sets random seed for numpy and both CPU/GPU pytorch routines.     If cuda is not available, this is silently ignored.     Note: numpy and pytorch do NOT use the same RNG. Given model class, instantiation arguments and weight file, load the model Performs a feed forward step of the     given model and returns model outputs given a torch.Tensor input. Performs a single step of backpropagation     and weight optimization Freeze or unfreeze weights Slice and normalise a batch of images for visualisation # Sentinel 1 # Sentinel 2 # gamma correction