From 2220de7cc0c88bede730c5a0ace6f3bf3c4ccd8c Mon Sep 17 00:00:00 2001
From: "Daniel.Frisinghelli" <daniel.frisinghelli@eurac.edu>
Date: Wed, 12 Aug 2020 17:32:42 +0200
Subject: [PATCH] Added enumerations of tested models, optimizers and loss
functions
---
pysegcnn/core/models.py | 13 +++++++++++++
1 file changed, 13 insertions(+)
diff --git a/pysegcnn/core/models.py b/pysegcnn/core/models.py
index cbf195b..6d840fc 100644
--- a/pysegcnn/core/models.py
+++ b/pysegcnn/core/models.py
@@ -7,11 +7,13 @@ Created on Fri Jun 26 16:31:36 2020
"""
# builtins
import os
+import enum
# externals
import numpy as np
import torch
import torch.nn as nn
+import torch.optim as optim
# locals
from pysegcnn.core.layers import (Encoder, Decoder, Conv2dPool, Conv2dUnpool,
@@ -139,3 +141,14 @@ class UNet(Network):
# classification
return self.classifier(x)
+
+
+class SupportedModels(enum.Enum):
+ Unet = UNet
+
+
+class SupportedOptimizers(enum.Enum):
+ Adam = optim.Adam
+
+class SupportedLossFunctions(enum.Enum):
+ CrossEntropy = nn.CrossEntropyLoss
--
GitLab