Merge pull request #84 from sony/feature/20241211-padim-on-cpu

CPU support for PaDiM and PatchCore
sony · Dec 13, 2024 · f20512d · f20512d
2 parents 227aa6f + f665a34
commit f20512d
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 8 deletions.
diff --git a/plugins/_Training/padim.py b/plugins/_Training/padim.py
@@ -32,8 +32,8 @@
 
 
 def get_model(model, batch_size, feature_ratio):
-    nn.set_default_context(get_extension_context('cudnn'))
     try:
+        nn.set_default_context(get_extension_context('cudnn'))
         x = nn.Variable()
         F.relu(x)
     except:
@@ -265,7 +265,7 @@ def func(args):
     # Create anomary detection model
     logger.log(99, 'Saving anomary detection model...')
     contents = {
-        'global_config': {'default_context': get_extension_context('cudnn')},
+        'global_config': {'default_context': nn.get_current_context()},
         'networks': [
             {'name': 'network',
              'batch_size': batch_size,

diff --git a/plugins/_Training/padim_c1.py b/plugins/_Training/padim_c1.py
@@ -35,8 +35,8 @@
 
 
 def get_model(model, batch_size, feature_ratio):
-    nn.set_default_context(get_extension_context('cudnn'))
     try:
+        nn.set_default_context(get_extension_context('cudnn'))
         x = nn.Variable()
         F.relu(x)
     except:
@@ -282,7 +282,7 @@ def func(args):
     # Create anomary detection model
     logger.log(99, 'Saving anomary detection model...')
     contents = {
-        'global_config': {'default_context': get_extension_context('cudnn')},
+        'global_config': {'default_context': nn.get_current_context()},
         'networks': [
             {'name': 'network',
              'batch_size': batch_size,

diff --git a/plugins/_Training/patchcore.py b/plugins/_Training/patchcore.py
@@ -33,8 +33,8 @@
 
 
 def get_model(model, batch_size, feature_ratio):
-    nn.set_default_context(get_extension_context('cudnn'))
     try:
+        nn.set_default_context(get_extension_context('cudnn'))
         x = nn.Variable()
         F.relu(x)
     except:
@@ -177,7 +177,7 @@ def func(args):
     # Create anomary detection model
     logger.log(99, 'Saving anomary detection model...')
     contents = {
-        'global_config': {'default_context': get_extension_context('cudnn')},
+        'global_config': {'default_context': nn.get_current_context()},
         'networks': [
             {'name': 'network',
              'batch_size': 1,

diff --git a/plugins/_Training/patchcore_c1.py b/plugins/_Training/patchcore_c1.py
@@ -27,15 +27,17 @@
 from nnabla import logger
 import nnabla_ext.cpu
 from nnabla.ext_utils import get_extension_context
+from nnabla.models.utils import get_model_home
 from nnabla.utils import nnabla_pb2
 from nnabla.utils.data_iterator import data_iterator_csv_dataset
+from nnabla.utils.download import download
 from nnabla.utils.progress import configure_progress, progress
 from nnabla.utils.save import save
 
 
 def get_model(model, batch_size, feature_ratio):
-    nn.set_default_context(get_extension_context('cudnn'))
     try:
+        nn.set_default_context(get_extension_context('cudnn'))
         x = nn.Variable()
         F.relu(x)
     except:
@@ -192,7 +194,7 @@ def func(args):
     # Create anomary detection model
     logger.log(99, 'Saving anomary detection model...')
     contents = {
-        'global_config': {'default_context': get_extension_context('cudnn')},
+        'global_config': {'default_context': nn.get_current_context()},
         'networks': [
             {'name': 'network',
              'batch_size': 1,