Update model.py

model.py

@@ -8,26 +8,21 @@ from torchvision.models import resnet18, resnet50
 from torchvision.models import ResNet18_Weights, ResNet50_Weights
 
 class DistMult(nn.Module):
-    def __init__(self, args, num_ent_uid, target_list, device, all_locs=None, num_habitat=None, all_timestamps=None):
+    def __init__(self, num_ent_uid, target_list, device, all_locs=None, num_habitat=None, all_timestamps=None):
         super(DistMult, self).__init__()
-        self.args = args
         self.num_ent_uid = num_ent_uid
 
         self.num_relations = 4
 
-        self.ent_embedding = torch.nn.Embedding(self.num_ent_uid, args.embedding_dim, sparse=False)
-        self.rel_embedding = torch.nn.Embedding(self.num_relations, args.embedding_dim, sparse=False)
+        self.ent_embedding = torch.nn.Embedding(self.num_ent_uid, 512, sparse=False)
+        self.rel_embedding = torch.nn.Embedding(self.num_relations, 512, sparse=False)
 
-        self.location_embedding = MLP(args.location_input_dim, args.embedding_dim, args.mlp_location_numlayer)
+        self.location_embedding = MLP(2, 512, 3)
 
-        self.time_embedding = MLP(args.time_input_dim, args.embedding_dim, args.mlp_time_numlayer)
+        self.time_embedding = MLP(1, 512, 3)
         
-        if self.args.img_embed_model == 'resnet50':
-            self.image_embedding = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
-            self.image_embedding.fc = nn.Linear(2048, args.embedding_dim)
-        else:
-            self.image_embedding = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
-            self.image_embedding.fc = nn.Linear(512, args.embedding_dim)
+        self.image_embedding = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+        self.image_embedding.fc = nn.Linear(2048, 512)
 
         self.target_list = target_list
         
@@ -36,7 +31,6 @@ class DistMult(nn.Module):
         if all_timestamps is not None:
             self.all_timestamps = all_timestamps.to(device)
 
-        self.args = args
         self.device = device
 
         self.init()