modelをddpmに固定

kenta-tsukaue · Jan 8, 2023 · f7e2f8c · f7e2f8c
1 parent 4a7a9d4
commit f7e2f8c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/models/ddpm.py b/models/ddpm.py
@@ -36,7 +36,6 @@
 default_initializer = layers.default_init
 
 
-@utils.register_model(name='ddpm')
 class DDPM(nn.Module):
   def __init__(self, config):
     super().__init__()

diff --git a/models/utils.py b/models/utils.py
@@ -18,6 +18,7 @@
 
 import torch
 import sde_lib
+from .ddpm import DDPM
 import numpy as np
 
 
@@ -87,8 +88,9 @@ def get_ddpm_params(config):
 
 def create_model(config):
   """Create the score model."""
-  model_name = config.model.name
-  score_model = get_model(model_name)(config)
+  print("モデルの名前は",config.model.name)
+  #score_model = get_model(model_name)(config)
+  score_model = DDPM(config)
   score_model = score_model.to(config.device)
   score_model = torch.nn.DataParallel(score_model)
   return score_model