Merge pull request #108 from kookmin-sw/mhsong-dev

Fix several errors

Author: MoohyunSong

automation/deploy_train/create_rayjob.py

@@ -256,7 +256,7 @@ def getTrainInfo(ModelClass, optimstr, lossstr, lr, model_dir):
       import torch.nn as nn
       model = ModelClass()
       if os.path.exists(f'{{model_dir}}/torch.pt'):
-        model.load_state_dict(torch.load("{{model_dir}}/torch.pt"))
+        model.load_state_dict(torch.load(f"{{model_dir}}/torch.pt"))
 
       optimizer = eval("optim."+optimstr+"(model.parameters(), lr=lr)")
       criterion = eval("nn."+lossstr+"()")
@@ -300,7 +300,7 @@ def train_func(config):
       test_loader = prepare_data_loader(test_loader)
 
       # 모델, 손실 함수 및 옵티마이저 설정
-      model, criterion, optimizer = getTrainInfo(ModelClass, optimstr=OPTIMIZER_STR, lossstr=LOSS_STR, lr=config["lr"], model_dir)
+      model, criterion, optimizer = getTrainInfo(ModelClass, optimstr=OPTIMIZER_STR, lossstr=LOSS_STR, lr=config["lr"], model_dir=model_dir)
       model = ray.train.torch.prepare_model(model)
       if torch.cuda.is_available():
         model = model.cuda()

inference/template_code/lambda_app.py

@@ -30,10 +30,29 @@
 
 try:
     model = ModelClass()
-    if os.path.exists('/tmp/model/torch.pt'):
-        model.load_state_dict(torch.load("/tmp/model/torch.pt", map_location=torch.device('cpu')))
+    model_path = '/tmp/model/torch.pt'
+
+    if os.path.exists(model_path):
+        try:
+            state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+            
+            # DataParallel로 래핑된 모델인지 확인
+            if 'module.' in list(state_dict.keys())[0]:
+                # 키에서 'module.'을 제거
+                from collections import OrderedDict
+                new_state_dict = OrderedDict()
+                for k, v in state_dict.items():
+                    new_state_dict[k.replace('module.', '')] = v
+                model.load_state_dict(new_state_dict)
+            else:
+                model.load_state_dict(state_dict)
+        except Exception as e:
+            print(f"모델 상태 딕셔너리를 로드하는 중 오류 발생: {e}")
+            os._exit(0)
+    else:
+        raise FileNotFoundError(f"모델 파일을 찾을 수 없습니다: {model_path}")
 except Exception as e:
-    print(f"Model load failed: {e}")
+    print(f"모델 로드 실패: {e}")
     os._exit(0)
 
 def handler(event, context):