cpu pls

models/model_main.py

@@ -5,6 +5,7 @@ from models.vgg_perceptual_loss import VGGPerceptualLoss
 from models.transformers import *
 from torch.autograd import Variable
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 class ModelMain(nn.Module):
 
     def __init__(self, opts, mode='train'):
@@ -72,7 +73,7 @@ class ModelMain(nn.Module):
         
         if mode in {'train', 'val'}:
             # seq decoding (training or val mode)
-            tgt_mask = Variable(subsequent_mask(self.opts.max_seq_len).type_as(ref_pad_mask.data)).unsqueeze(0).expand(z.size(0), -1, -1, -1).cuda().float()
+            tgt_mask = Variable(subsequent_mask(self.opts.max_seq_len).type_as(ref_pad_mask.data)).unsqueeze(0).expand(z.size(0), -1, -1, -1).to(device).float()
             command_logits, args_logits, attn = self.transformer_seqdec(x=trg_seq_shifted, memory=latent_feat_seq, trg_char=trg_cls, tgt_mask=tgt_mask)
             command_logits_2, args_logits_2 = self.transformer_seqdec.parallel_decoder(command_logits, args_logits, memory=latent_feat_seq.detach(), trg_char=trg_cls)
 
@@ -97,10 +98,10 @@ class ModelMain(nn.Module):
         else: # testing (inference)
 
             trg_len = trg_seq_shifted.size(0)
-            sampled_svg = torch.zeros(1, trg_seq.size(1), self.opts.dim_seq_short).cuda()
+            sampled_svg = torch.zeros(1, trg_seq.size(1), self.opts.dim_seq_short).to(device)
 
             for t in range(0, trg_len):
-                tgt_mask = Variable(subsequent_mask(sampled_svg.size(0)).type_as(ref_seq_cat.data)).unsqueeze(0).expand(sampled_svg.size(1), -1, -1, -1).cuda().float()
+                tgt_mask = Variable(subsequent_mask(sampled_svg.size(0)).type_as(ref_seq_cat.data)).unsqueeze(0).expand(sampled_svg.size(1), -1, -1, -1).to(device).float()
                 command_logits, args_logits, attn = self.transformer_seqdec(x=sampled_svg, memory=latent_feat_seq, trg_char=trg_cls, tgt_mask=tgt_mask)
                 prob_comand = F.softmax(command_logits[:, -1, :], -1)
                 prob_args = F.softmax(args_logits[:, -1, :], -1)
@@ -151,29 +152,29 @@ class ModelMain(nn.Module):
 
 
         if mode == 'train':
-            ref_cls = torch.randint(0, self.opts.char_num, (input_image.size(0), self.opts.ref_nshot)).cuda()
+            ref_cls = torch.randint(0, self.opts.char_num, (input_image.size(0), self.opts.ref_nshot)).to(device)
             if opts.ref_nshot == 52: # For ENG to TH
-                ref_cls_upper = torch.randint(0, 26, (input_image.size(0), self.opts.ref_nshot // 2)).cuda()
-                ref_cls_lower = torch.randint(26, 52, (input_image.size(0), self.opts.ref_nshot // 2)).cuda()
+                ref_cls_upper = torch.randint(0, 26, (input_image.size(0), self.opts.ref_nshot // 2)).to(device)
+                ref_cls_lower = torch.randint(26, 52, (input_image.size(0), self.opts.ref_nshot // 2)).to(device)
                 ref_cls = torch.cat((ref_cls_upper, ref_cls_lower), -1)
         elif mode == 'val':
-            ref_cls = torch.arange(0, self.opts.ref_nshot, 1).cuda().unsqueeze(0).expand(input_image.size(0), -1)
+            ref_cls = torch.arange(0, self.opts.ref_nshot, 1).to(device).unsqueeze(0).expand(input_image.size(0), -1)
         else:
             ref_ids = self.opts.ref_char_ids.split(',')
             ref_ids = list(map(int, ref_ids))
             assert len(ref_ids) == self.opts.ref_nshot
-            ref_cls = torch.tensor(ref_ids).cuda().unsqueeze(0).expand(self.opts.char_num, -1)
+            ref_cls = torch.tensor(ref_ids).to(device).unsqueeze(0).expand(self.opts.char_num, -1)
         
         
         
         if mode in {'train', 'val'}:
-            trg_cls = torch.randint(0, self.opts.char_num, (input_image.size(0), 1)).cuda()
+            trg_cls = torch.randint(0, self.opts.char_num, (input_image.size(0), 1)).to(device)
             if opts.ref_nshot == 52:
-                trg_cls = torch.randint(52, opts.char_num, (input_image.size(0), 1)).cuda()
+                trg_cls = torch.randint(52, opts.char_num, (input_image.size(0), 1)).to(device)
         else:
-            trg_cls = torch.arange(0, self.opts.char_num).cuda()
+            trg_cls = torch.arange(0, self.opts.char_num).to(device)
             if opts.ref_nshot == 52:
-                trg_cls = torch.randint(52, opts.char_num, (input_image.size(0), 1)).cuda()
+                trg_cls = torch.randint(52, opts.char_num, (input_image.size(0), 1)).to(device)
             trg_cls = trg_cls.view(self.opts.char_num, 1)
             input_image = input_image.expand(self.opts.char_num, -1, -1, -1)
             input_sequence = input_sequence.expand(self.opts.char_num, -1, -1, -1)
@@ -205,7 +206,7 @@ class ModelMain(nn.Module):
         ref_pad_mask = torch.zeros(ref_seqlen_cat.size(0), self.opts.max_seq_len) # value = 1 means pos to be masked
         for i in range(ref_seqlen_cat.size(0)):
             ref_pad_mask[i,:ref_seqlen_cat[i]] = 1
-        ref_pad_mask = ref_pad_mask.cuda().float().unsqueeze(1)
+        ref_pad_mask = ref_pad_mask.to(device).float().unsqueeze(1)
         trg_seqlen = util_funcs.select_seqlens(input_seqlen, trg_cls, self.opts)
         trg_seqlen = trg_seqlen.squeeze()
 

test_few_shot.py

@@ -23,16 +23,18 @@ def test_main_model(opts):
         dir_res = os.path.join(f"{opts.exp_path}", "experiments/", opts.name_exp, "results")
 
     test_loader = get_loader(opts.data_root, opts.img_size, opts.language, opts.char_num, opts.max_seq_len, opts.dim_seq, opts.batch_size, 'test')
-    if torch.cuda.is_available():
-        device = torch.device("cuda")
-    else:
-        device = torch.device("cpu")
+    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print("Inference With Device:", device)
     if opts.streamlit:
         st.write("Loading Model Weight...")
+        st.write("Inference With Device:", device)
+
     model_main = ModelMain(opts)
     path_ckpt = os.path.join(f"{opts.model_path}")
-    model_main.load_state_dict(torch.load(path_ckpt, map_location=torch.device('cpu'))['model'])
+    model_main.load_state_dict(torch.load(path_ckpt, map_location=device)['model'])
     model_main.to(device)
+    
     model_main.eval()
     with torch.no_grad():