Update models/monet.py

models/monet.py

@@ -69,7 +69,7 @@ class MAL(nn.Module):
         Multi-view Attention Learning (MAL) module
     """
 
-    def __init__(self, in_dim=768, feature_num=4, feature_size=28):
+    def __init__(self, in_dim=768, feature_num=4, feature_size=28, is_gpu=True):
         super().__init__()
 
         self.channel_attention = Attention_Block(in_dim * feature_num)  # Channel-wise self attention
@@ -82,9 +82,14 @@ class MAL(nn.Module):
 
         self.feature_num = feature_num
         self.in_dim = in_dim
+        self.is_gpu = is_gpu
 
     def forward(self, features):
-        feature = torch.tensor([]).cuda()
+        if self.is_gpu:
+            feature = torch.tensor([]).cuda()
+        else:
+            feature = torch.tensor([])
+
         for index, _ in enumerate(features):
             feature = torch.cat((feature, self.attention_module[index](features[index]).unsqueeze(0)), dim=0)
         features = feature
@@ -118,7 +123,7 @@ class SaveOutput:
 
 
 class MoNet(nn.Module):
-    def __init__(self, config, patch_size=8, drop=0.1, dim_mlp=768, img_size=224):
+    def __init__(self, config, patch_size=8, drop=0.1, dim_mlp=768, img_size=224, is_gpu=True):
         super().__init__()
         self.img_size = img_size
         self.input_size = img_size // patch_size
@@ -136,10 +141,10 @@ class MoNet(nn.Module):
 
         self.MALs = nn.ModuleList()
         for _ in range(config.mal_num):
-            self.MALs.append(MAL())
+            self.MALs.append(MAL(is_gpu=is_gpu))
 
         # Image Quality Score Regression
-        self.fusion_wam = MAL(feature_num=config.mal_num)
+        self.fusion_wam = MAL(feature_num=config.mal_num, is_gpu=is_gpu)
         self.block = Block(dim_mlp, 12)
         self.cnn = nn.Sequential(
             nn.Conv2d(dim_mlp, 256, 5),
@@ -163,6 +168,8 @@ class MoNet(nn.Module):
             nn.Sigmoid()
         )
 
+        self.is_gpu = is_gpu
+
     def extract_feature(self, save_output, block_index=[2, 5, 8, 11]):
         x1 = save_output.outputs[block_index[0]][:, 1:]
         x2 = save_output.outputs[block_index[1]][:, 1:]
@@ -182,7 +189,11 @@ class MoNet(nn.Module):
         x = x.permute(1, 0, 2, 3, 4)  # bs, 4, 768, 28 * 28
 
         # Different Opinion Features (DOF)
-        DOF = torch.tensor([]).cuda()
+        if self.is_gpu:
+            DOF = torch.tensor([]).cuda()
+        else:
+            DOF = torch.tensor([])
+        
         for index, _ in enumerate(self.MALs):
             DOF = torch.cat((DOF, self.MALs[index](x).unsqueeze(0)), dim=0)
         DOF = rearrange(DOF, 'n c d (w h) -> n c d w h', w=self.input_size, h=self.input_size)  # 3, bs, 768, 28, 28