cached models improve speed

app.py

@@ -8,38 +8,41 @@ from matplotlib import colors
 
 if not hasattr(st, 'paths'):
     st.paths = None
+if not hasattr(st, 'daily_model'):
+    best_model_daily_file_name = "best_model_daily.pth"
+    best_model_annual_file_name = "best_model_annual.pth"
+
+    first_input_batch = torch.zeros(71, 9, 5, 48, 48)
+    # first_input_batch = first_input_batch.view(-1, *first_input_batch.shape[2:])
+    st.daily_model = FPN(opt, first_input_batch, opt.win_size)
+    st.annual_model = SimpleNN(opt)
+
+    if torch.cuda.is_available():
+        st.daily_model = torch.nn.DataParallel(st.daily_model).cuda()
+        st.annual_model = torch.nn.DataParallel(st.annual_model).cuda()
+        st.daily_model = torch.nn.DataParallel(st.daily_model).cuda()
+        st.annual_model = torch.nn.DataParallel(st.annual_model).cuda()
+    else:
+        st.daily_model = torch.nn.DataParallel(st.daily_model).cpu()
+        st.annual_model = torch.nn.DataParallel(st.annual_model).cpu()
+        st.daily_model = torch.nn.DataParallel(st.daily_model).cpu()
+        st.annual_model = torch.nn.DataParallel(st.annual_model).cpu()
+
+    print('trying to resume previous saved models...')
+    state = resume(
+        os.path.join(opt.resume_path, best_model_daily_file_name),
+        model=st.daily_model, optimizer=None)
+    state = resume(
+        os.path.join(opt.resume_path, best_model_annual_file_name),
+        model=st.annual_model, optimizer=None)
+    st.daily_model = st.daily_model.eval()
+    st.annual_model = st.annual_model.eval()
+
 
 # Load Model
 # @title Load pretrained weights
 
-best_model_daily_file_name = "best_model_daily.pth"
-best_model_annual_file_name = "best_model_annual.pth"
-
-first_input_batch = torch.zeros(71, 9, 5, 48, 48)
-# first_input_batch = first_input_batch.view(-1, *first_input_batch.shape[2:])
-daily_model = FPN(opt, first_input_batch, opt.win_size)
-annual_model = SimpleNN(opt)
-
-if torch.cuda.is_available():
-    daily_model = torch.nn.DataParallel(daily_model).cuda()
-    annual_model = torch.nn.DataParallel(annual_model).cuda()
-    daily_model = torch.nn.DataParallel(daily_model).cuda()
-    annual_model = torch.nn.DataParallel(annual_model).cuda()
-else:
-    daily_model = torch.nn.DataParallel(daily_model).cpu()
-    annual_model = torch.nn.DataParallel(annual_model).cpu()
-    daily_model = torch.nn.DataParallel(daily_model).cpu()
-    annual_model = torch.nn.DataParallel(annual_model).cpu()
-
-print('trying to resume previous saved models...')
-state = resume(
-    os.path.join(opt.resume_path, best_model_daily_file_name),
-    model=daily_model, optimizer=None)
-state = resume(
-    os.path.join(opt.resume_path, best_model_annual_file_name),
-    model=annual_model, optimizer=None)
-daily_model = daily_model.eval()
-annual_model = annual_model.eval()
+
 
 st.title('Lombardia Sentinel 2 daily Crop Mapping')
 st.markdown('Using a daily FPN and giving a zip that contains 30 tiff with 7 channels, correctly named you can reach prediction of crop mapping og the area.')
@@ -85,14 +88,14 @@ if sample_path is not None:
             if torch.cuda.is_available():
                 x_dailies = x_dailies.cuda()
 
-            feat_daily, outs_daily = daily_model.forward(x_dailies)
+            feat_daily, outs_daily = st.daily_model.forward(x_dailies)
             # return to original size of batch and year
             outs_daily = outs_daily.view(
                 opt.batch_size, opt.sample_duration, *outs_daily.shape[1:])
             feat_daily = feat_daily.view(
                 opt.batch_size, opt.sample_duration, *feat_daily.shape[1:])
 
-            _, out_annual = annual_model.forward(feat_daily)
+            _, out_annual = st.annual_model.forward(feat_daily)
             pred_annual = torch.argmax(out_annual, dim=1).squeeze(1)
             pred_annual = pred_annual.cpu().numpy()
             # Remapping the labels
@@ -158,7 +161,7 @@ if st.paths is not None:
                                st.paths, index=st.paths.index('patch-pred-nn.tif'))
 
     file_path = os.path.join(folder, file_picker)
-    print(file_path)
+    # print(file_path)
     target, profile = read(file_path)
     target = np.squeeze(target)
     target = [classes_color_map[p] for p in target]
@@ -169,7 +172,7 @@ if st.paths is not None:
 
     markdown_legend = ''
     for c, l in zip(classes_color_map, labels_map):
-        print(colors.to_hex(c))
+        # print(colors.to_hex(c))
         markdown_legend += f'<div style="color:gray;background-color: {colors.to_hex(c)};">{l}</div><br>'
 
     col1, col2 = st.columns(2)