nathangavenski / il-datasets Goto Github PK

def accuracy(prediction: Tensor, ground_truth: Tensor) -> Number:
    """Compute the accuracy for a model. The accuracy returned is the percentage from 0 to 100.

    Args:
        prediction (torch.Tensor): logits from a model.
        ground_truth (torch.Tensor): ground truth class.

    Raises:
        ValueError: if predictions and ground_truth are not torch.Tensor.
        ValueError: if predictions are not two dimensional.
        ValueError: if ground_truth is not one dimensional.

    Returns:
        accuracy (Number): accuracy between 0 and 100 for a model.
    """

    if not isinstance(prediction, Tensor) or not isinstance(ground_truth, Tensor):
        raise ValueError("'prediction' and 'ground truth' should be a tensor")

    if len(prediction.size()) != 2:
        raise ValueError("'prediction' and 'ground truth' need to be 2 dimensional.")

    if len(ground_truth.size()) != 1:
        raise ValueError("'ground truth' need to be 1 dimensional.")

    return ((argmax(prediction, 1) == ground_truth).sum().item() / ground_truth.size(0)) * 100

class Method(ABC):
    """Base class for all methods."""

    __version__ = "1.0.0"
    __author__ = "Nathan Gavenski"
    __method_name__ = "Abstract Method"

    def __init__(
        self,
        environment: Env,
        environment_parameters: Dict[str, Any],
        discrete_loss: nn.Module = nn.CrossEntropyLoss,
        continuous_loss: nn.Module = nn.MSELoss,
        optimizer_fn: optim.Optimizer = optim.Adam,
    ) -> None:
        """Initialize base class."""
        super().__init__()
        self.environment = environment
        self.discrete = isinstance(environment.action_space, spaces.Discrete)
        self.discrete |= isinstance(environment.action_space, gym_spaces.Discrete)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        observation_size = environment.observation_space.shape[0]

        if self.discrete:
            action_size = environment.action_space.n
            self.loss_fn = discrete_loss()
        else:
            action_size = environment.action_space.shape[0]
            self.loss_fn = continuous_loss()

        self.policy = MLP(observation_size, action_size)

        self.optimizer_fn = optimizer_fn(
            self.policy.parameters(),
            **environment_parameters
        )

def step(
      self,
      action: Union[float, int]
) -> Union[
  Tuple[List[float], float, bool, bool, Dict[str, Any]],
  Tuple[List[float], float, bool, Dict[str, Any]]
]:
  """
  Perform an action in the environment and return the appropriate return
  according to version.
  """
  gym_return = self.env.step(action)
  if self.version == "newest":
      state, reward, terminated, truncated, info = gym_return
      return state, reward, terminated or truncated, info

  return gym_return