deepliif 1.1.6__py3-none-any.whl → 1.1.8__py3-none-any.whl

deepliif/models/__init__.py

@@ -31,15 +31,32 @@ import numpy as np
 from dask import delayed, compute
 
 from deepliif.util import *
-from deepliif.util.util import tensor_to_pil
+from deepliif.util.util import tensor_to_pil, check_multi_scale
 from deepliif.data import transform
-from deepliif.postprocessing import adjust_marker, adjust_dapi, compute_IHC_scoring, \
-    overlay_final_segmentation_mask, create_final_segmentation_mask_with_boundaries, create_basic_segmentation_mask
+from deepliif.postprocessing import compute_results
+from deepliif.options import Options, print_options
 
 from .base_model import BaseModel
+
+# import for init purpose, not used in this script
 from .DeepLIIF_model import DeepLIIFModel
-from .networks import get_norm_layer, ResnetGenerator, UnetGenerator
+from .DeepLIIFExt_model import DeepLIIFExtModel
+
 
+@lru_cache
+def get_opt(model_dir, mode='test'):
+    """
+    mode: test or train, currently only functions used for inference utilize get_opt so it
+          defaults to test
+    """
+    if mode == 'train':
+        opt = Options(path_file=os.path.join(model_dir,'train_opt.txt'), mode=mode)
+    elif mode == 'test':
+        try:
+            opt = Options(path_file=os.path.join(model_dir,'test_opt.txt'), mode=mode)
+        except:
+            opt = Options(path_file=os.path.join(model_dir,'train_opt.txt'), mode=mode)
+    return opt
 
 def find_model_using_name(model_name):
     """Import the module "models/[model_name]_model.py".
@@ -94,87 +111,75 @@ def load_torchscript_model(model_pt_path, device):
     return net
 
 
-def read_model_params(file_addr):
-    with open(file_addr) as f:
-        lines = f.readlines()
-    param_dict = {}
-    for line in lines:
-        if ':' in line:
-            key = line.split(':')[0].strip()
-            val = line.split(':')[1].split('[')[0].strip()
-            param_dict[key] = val
-    print(param_dict)
-    return param_dict
-
-
-def load_eager_models(model_dir, devices):
-    input_nc = 3
-    output_nc = 3
-    ngf = 64
-    norm = 'batch'
-    use_dropout = True
-    padding_type = 'zero'
-
-    files = os.listdir(model_dir)
-    for f in files:
-        if 'train_opt.txt' in f:
-            param_dict = read_model_params(os.path.join(model_dir, f))
-            input_nc = int(param_dict['input_nc'])
-            output_nc = int(param_dict['output_nc'])
-            ngf = int(param_dict['ngf'])
-            norm = param_dict['norm']
-            use_dropout = False if param_dict['no_dropout'] == 'True' else True
-            padding_type = param_dict['padding']
-
-    norm_layer = get_norm_layer(norm_type=norm)
 
-    nets = {}
-    for n in ['G1', 'G2', 'G3', 'G4']:
-        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9, padding_type=padding_type)
-        net.load_state_dict(torch.load(
-            os.path.join(model_dir, f'latest_net_{n}.pth'),
-            map_location=devices[n]
-        ))
-        nets[n] = disable_batchnorm_tracking_stats(net)
-        nets[n].eval()
-
-    for n in ['G51', 'G52', 'G53', 'G54', 'G55']:
-        net = UnetGenerator(input_nc, output_nc, 9, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
-        net.load_state_dict(torch.load(
-            os.path.join(model_dir, f'latest_net_{n}.pth'),
-            map_location=devices[n]
-        ))
-        nets[n] = disable_batchnorm_tracking_stats(net)
-        nets[n].eval()
+def load_eager_models(opt, devices):
+    # create a model given model and other options
+    model = create_model(opt)
+    # regular setup: load and print networks; create schedulers
+    model.setup(opt)
 
+    nets = {}
+    for name in model.model_names:
+        if isinstance(name, str):
+            if '_' in name:
+                net = getattr(model, 'net' + name.split('_')[0])[int(name.split('_')[-1]) - 1]
+            else:
+                net = getattr(model, 'net' + name)
+    
+            if opt.phase != 'train':
+                net.eval()
+                net = disable_batchnorm_tracking_stats(net)
+            
+            nets[name] = net
+            nets[name].to(devices[name])
+            
     return nets
 
 
 @lru_cache
-def init_nets(model_dir, eager_mode=False):
+def init_nets(model_dir, eager_mode=False, opt=None, phase='test'):
     """
     Init DeepLIIF networks so that every net in
     the same group is deployed on the same GPU
-    """
-    net_groups = [
-        ('G1', 'G52'),
-        ('G2', 'G53'),
-        ('G3', 'G54'),
-        ('G4', 'G55'),
-        ('G51',)
-    ]
-
-    number_of_gpus = torch.cuda.device_count()
-    # number_of_gpus = 0
-    if number_of_gpus:
+    
+    opt_args: to overwrite opt arguments in train_opt.txt, typically used in inference stage
+              for example, opt_args={'phase':'test'}
+    """ 
+    if opt is None:
+        opt = get_opt(model_dir, mode=phase)
+        opt.use_dp = False
+        print_options(opt)
+    
+    if opt.model == 'DeepLIIF':
+        net_groups = [
+            ('G1', 'G52'),
+            ('G2', 'G53'),
+            ('G3', 'G54'),
+            ('G4', 'G55'),
+            ('G51',)
+        ]
+    elif opt.model == 'DeepLIIFExt':
+        if opt.seg_gen:
+            net_groups = [(f'G_{i+1}',f'GS_{i+1}') for i in range(opt.modalities_no)]
+        else:
+            net_groups = [(f'G_{i+1}',) for i in range(opt.modalities_no)]
+    else:
+        raise Exception(f'init_nets() not implemented for model {opt.model}')
+
+    number_of_gpus_all = torch.cuda.device_count()
+    number_of_gpus = len(opt.gpu_ids)
+    print(number_of_gpus)
+    if number_of_gpus > 0:
+        mapping_gpu_ids = {i:idx for i,idx in enumerate(opt.gpu_ids)}
         chunks = [itertools.chain.from_iterable(c) for c in chunker(net_groups, number_of_gpus)]
         # chunks = chunks[1:]
-        devices = {n: torch.device(f'cuda:{i}') for i, g in enumerate(chunks) for n in g}
+        devices = {n: torch.device(f'cuda:{mapping_gpu_ids[i]}') for i, g in enumerate(chunks) for n in g}
+        # devices = {n: torch.device(f'cuda:{i}') for i, g in enumerate(chunks) for n in g}
     else:
         devices = {n: torch.device('cpu') for n in itertools.chain.from_iterable(net_groups)}
 
     if eager_mode:
-        return load_eager_models(model_dir, devices)
+        return load_eager_models(opt, devices)
 
     return {
         n: load_torchscript_model(os.path.join(model_dir, f'{n}.pt'), device=d)
@@ -190,14 +195,18 @@ def compute_overlap(img_size, tile_size):
     return tile_size // 4
 
 
-def run_torchserve(img, model_path=None, eager_mode=False):
+def run_torchserve(img, model_path=None, eager_mode=False, opt=None):
     """
     eager_mode: not used in this function; put in place to be consistent with run_dask
            so that run_wrapper() could call either this function or run_dask with
            same syntax
+    opt: same as eager_mode
     """
     buffer = BytesIO()
-    torch.save(transform(img.resize((512, 512))), buffer)
+    if opt.model == 'DeepLIIFExt':
+        torch.save(transform(img.resize((1024, 1024))), buffer)
+    else:
+        torch.save(transform(img.resize((512, 512))), buffer)
 
     torchserve_host = os.getenv('TORCHSERVE_HOST', 'http://localhost')
     res = requests.post(
@@ -213,33 +222,55 @@ def run_torchserve(img, model_path=None, eager_mode=False):
     return {k: tensor_to_pil(deserialize_tensor(v)) for k, v in res.json().items()}
 
 
-def run_dask(img, model_path, eager_mode=False):
+def run_dask(img, model_path, eager_mode=False, opt=None):
     model_dir = os.getenv('DEEPLIIF_MODEL_DIR', model_path)
-    nets = init_nets(model_dir, eager_mode)
-
-    ts = transform(img.resize((512, 512)))
+    nets = init_nets(model_dir, eager_mode, opt)
+    
+    if opt.model == 'DeepLIIFExt':
+        ts = transform(img.resize((1024, 1024)))
+    else:
+        ts = transform(img.resize((512, 512)))
 
     @delayed
     def forward(input, model):
         with torch.no_grad():
             return model(input.to(next(model.parameters()).device))
+    
+    if opt.model == 'DeepLIIF':
+        seg_map = {'G1': 'G52', 'G2': 'G53', 'G3': 'G54', 'G4': 'G55'}
+        
+        lazy_gens = {k: forward(ts, nets[k]) for k in seg_map}
+        gens = compute(lazy_gens)[0]
+        
+        lazy_segs = {v: forward(gens[k], nets[v]).to(torch.device('cpu')) for k, v in seg_map.items()}
+        lazy_segs['G51'] = forward(ts, nets['G51']).to(torch.device('cpu'))
+        segs = compute(lazy_segs)[0]
+    
+        seg_weights = [0.25, 0.25, 0.25, 0, 0.25]
+        seg = torch.stack([torch.mul(n, w) for n, w in zip(segs.values(), seg_weights)]).sum(dim=0)
+    
+        res = {k: tensor_to_pil(v) for k, v in gens.items()}
+        res['G5'] = tensor_to_pil(seg)
+    
+        return res
+    elif opt.model == 'DeepLIIFExt':
+        seg_map = {'G_' + str(i): 'GS_' + str(i) for i in range(1, opt.modalities_no + 1)}
+        
+        lazy_gens = {k: forward(ts, nets[k]) for k in seg_map}
+        gens = compute(lazy_gens)[0]
+        
+        res = {k: tensor_to_pil(v) for k, v in gens.items()}
+    
+        if opt.seg_gen:
+            lazy_segs = {v: forward(torch.cat([ts.to(torch.device('cpu')), gens[next(iter(seg_map))].to(torch.device('cpu')), gens[k].to(torch.device('cpu'))], 1), nets[v]).to(torch.device('cpu')) for k, v in seg_map.items()}
+            segs = compute(lazy_segs)[0]
+            res.update({k: tensor_to_pil(v) for k, v in segs.items()})
+    
+        return res
+    else:
+        raise Exception(f'run_dask() not implemented for {opt.model}')
 
-    seg_map = {'G1': 'G52', 'G2': 'G53', 'G3': 'G54', 'G4': 'G55'}
-
-    lazy_gens = {k: forward(ts, nets[k]) for k in seg_map}
-    gens = compute(lazy_gens)[0]
-
-    lazy_segs = {v: forward(gens[k], nets[v]).to(torch.device('cpu')) for k, v in seg_map.items()}
-    lazy_segs['G51'] = forward(ts, nets['G51']).to(torch.device('cpu'))
-    segs = compute(lazy_segs)[0]
-
-    seg_weights = [0.25, 0.25, 0.25, 0, 0.25]
-    seg = torch.stack([torch.mul(n, w) for n, w in zip(segs.values(), seg_weights)]).sum(dim=0)
-
-    res = {k: tensor_to_pil(v) for k, v in gens.items()}
-    res['G5'] = tensor_to_pil(seg)
-
-    return res
+    
 
 
 def is_empty(tile):
@@ -247,17 +278,27 @@ def is_empty(tile):
     return True if calculate_background_area(tile) > 98 else False
 
 
-def run_wrapper(tile, run_fn, model_path, eager_mode=False):
-    if is_empty(tile):
-        return {
-            'G1': Image.new(mode='RGB', size=(512, 512), color=(201, 211, 208)),
-            'G2': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
-            'G3': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
-            'G4': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
-            'G5': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0))
-        }
+def run_wrapper(tile, run_fn, model_path, eager_mode=False, opt=None):
+    if opt.model == 'DeepLIIF':
+        if is_empty(tile):
+            return {
+                'G1': Image.new(mode='RGB', size=(512, 512), color=(201, 211, 208)),
+                'G2': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
+                'G3': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                'G4': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
+                'G5': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0))
+            }
+        else:
+            return run_fn(tile, model_path, eager_mode, opt)
+    elif opt.model == 'DeepLIIFExt':
+        if is_empty(tile):
+            res = {'G_' + str(i): Image.new(mode='RGB', size=(512, 512)) for i in range(1, opt.modalities_no + 1)}
+            res.update({'GS_' + str(i): Image.new(mode='RGB', size=(512, 512)) for i in range(1, opt.modalities_no + 1)})
+            return res
+        else:
+            return run_fn(tile, model_path, eager_mode, opt)
     else:
-        return run_fn(tile, model_path, eager_mode)
+        raise Exception(f'run_wrapper() not implemented for model {opt.model}')
 
 
 def inference_old(img, tile_size, overlap_size, model_path, use_torchserve=False, eager_mode=False,
@@ -307,71 +348,115 @@ def inference_old(img, tile_size, overlap_size, model_path, use_torchserve=False
 
 
 def inference(img, tile_size, overlap_size, model_path, use_torchserve=False, eager_mode=False,
-              color_dapi=False, color_marker=False):
-
-    rescaled, rows, cols = format_image_for_tiling(img, tile_size, overlap_size)
-
-    run_fn = run_torchserve if use_torchserve else run_dask
-
-    images = {}
-    images['Hema'] = create_image_for_stitching(tile_size, rows, cols)
-    images['DAPI'] = create_image_for_stitching(tile_size, rows, cols)
-    images['Lap2'] = create_image_for_stitching(tile_size, rows, cols)
-    images['Marker'] = create_image_for_stitching(tile_size, rows, cols)
-    images['Seg'] = create_image_for_stitching(tile_size, rows, cols)
-
-    for i in range(cols):
-        for j in range(rows):
-            tile = extract_tile(rescaled, tile_size, overlap_size, i, j)
-            res = run_fn(tile, model_path, eager_mode)
-
-            stitch_tile(images['Hema'], res['G1'], tile_size, overlap_size, i, j)
-            stitch_tile(images['DAPI'], res['G2'], tile_size, overlap_size, i, j)
-            stitch_tile(images['Lap2'], res['G3'], tile_size, overlap_size, i, j)
-            stitch_tile(images['Marker'], res['G4'], tile_size, overlap_size, i, j)
-            stitch_tile(images['Seg'], res['G5'], tile_size, overlap_size, i, j)
-
-    images['Hema'] = images['Hema'].resize(img.size)
-    images['DAPI'] = images['DAPI'].resize(img.size)
-    images['Lap2'] = images['Lap2'].resize(img.size)
-    images['Marker'] = images['Marker'].resize(img.size)
-    images['Seg'] = images['Seg'].resize(img.size)
-
-    if color_dapi:
-        matrix = (       0,        0,        0, 0,
-                  299/1000, 587/1000, 114/1000, 0,
-                  299/1000, 587/1000, 114/1000, 0)
-        images['DAPI'] = images['DAPI'].convert('RGB', matrix)
-
-    if color_marker:
-        matrix = (299/1000, 587/1000, 114/1000, 0,
-                  299/1000, 587/1000, 114/1000, 0,
-                         0,        0,        0, 0)
-        images['Marker'] = images['Marker'].convert('RGB', matrix)
-
-    return images
-
-
-def postprocess(img, seg_img, thresh=80, noise_objects_size=20, small_object_size=50):
-    mask_image = create_basic_segmentation_mask(np.array(img), np.array(seg_img),
-                                                thresh, noise_objects_size, small_object_size)
-    images = {}
-    images['SegOverlaid'] = Image.fromarray(overlay_final_segmentation_mask(np.array(img), mask_image))
-    images['SegRefined'] = Image.fromarray(create_final_segmentation_mask_with_boundaries(np.array(mask_image)))
-
-    all_cells_no, positive_cells_no, negative_cells_no, IHC_score = compute_IHC_scoring(mask_image)
-    scoring = {
-        'num_total': all_cells_no,
-        'num_pos': positive_cells_no,
-        'num_neg': negative_cells_no,
-        'percent_pos': IHC_score
-    }
+              color_dapi=False, color_marker=False, opt=None):
+    if not opt:
+        opt = get_opt(model_path)
+        print_options(opt)
+    
+    if opt.model == 'DeepLIIF':
+        rescaled, rows, cols = format_image_for_tiling(img, tile_size, overlap_size)
+    
+        run_fn = run_torchserve if use_torchserve else run_dask
+    
+        images = {}
+        images['Hema'] = create_image_for_stitching(tile_size, rows, cols)
+        images['DAPI'] = create_image_for_stitching(tile_size, rows, cols)
+        images['Lap2'] = create_image_for_stitching(tile_size, rows, cols)
+        images['Marker'] = create_image_for_stitching(tile_size, rows, cols)
+        images['Seg'] = create_image_for_stitching(tile_size, rows, cols)
+    
+        for i in range(cols):
+            for j in range(rows):
+                tile = extract_tile(rescaled, tile_size, overlap_size, i, j)
+                res = run_wrapper(tile, run_fn, model_path, eager_mode, opt)
+    
+                stitch_tile(images['Hema'], res['G1'], tile_size, overlap_size, i, j)
+                stitch_tile(images['DAPI'], res['G2'], tile_size, overlap_size, i, j)
+                stitch_tile(images['Lap2'], res['G3'], tile_size, overlap_size, i, j)
+                stitch_tile(images['Marker'], res['G4'], tile_size, overlap_size, i, j)
+                stitch_tile(images['Seg'], res['G5'], tile_size, overlap_size, i, j)
+    
+        images['Hema'] = images['Hema'].resize(img.size)
+        images['DAPI'] = images['DAPI'].resize(img.size)
+        images['Lap2'] = images['Lap2'].resize(img.size)
+        images['Marker'] = images['Marker'].resize(img.size)
+        images['Seg'] = images['Seg'].resize(img.size)
+    
+        if color_dapi:
+            matrix = (       0,        0,        0, 0,
+                      299/1000, 587/1000, 114/1000, 0,
+                      299/1000, 587/1000, 114/1000, 0)
+            images['DAPI'] = images['DAPI'].convert('RGB', matrix)
+    
+        if color_marker:
+            matrix = (299/1000, 587/1000, 114/1000, 0,
+                      299/1000, 587/1000, 114/1000, 0,
+                             0,        0,        0, 0)
+            images['Marker'] = images['Marker'].convert('RGB', matrix)
+    
+        return images
+        
+    elif opt.model == 'DeepLIIFExt':
+        #param_dict = read_train_options(model_path)
+        #modalities_no = int(param_dict['modalities_no']) if param_dict else 4
+        #seg_gen = (param_dict['seg_gen'] == 'True') if param_dict else True
+    
+        tiles = list(generate_tiles(img, tile_size, overlap_size))
+    
+        run_fn = run_torchserve if use_torchserve else run_dask
+        res = [Tile(t.i, t.j, run_wrapper(t.img, run_fn, model_path, eager_mode, opt)) for t in tiles]
+    
+        def get_net_tiles(n):
+            return [Tile(t.i, t.j, t.img[n]) for t in res]
+    
+        images = {}
+    
+        for i in range(1, opt.modalities_no + 1):
+            images['mod' + str(i)] = stitch(get_net_tiles('G_' + str(i)), tile_size, overlap_size).resize(img.size)
+    
+        if opt.seg_gen:
+            for i in range(1, opt.modalities_no + 1):
+                images['Seg' + str(i)] = stitch(get_net_tiles('GS_' + str(i)), tile_size, overlap_size).resize(img.size)
+    
+        return images
+    
+    else:
+        raise Exception(f'inference() not implemented for model {opt.model}')
+
+
+def postprocess(orig, images, tile_size, seg_thresh=150, size_thresh='default', marker_thresh='default', size_thresh_upper=None, opt=None):
+    if opt.model == 'DeepLIIF':
+        resolution = '40x' if tile_size > 384 else ('20x' if tile_size > 192 else '10x')
+        overlay, refined, scoring = compute_results(np.array(orig), np.array(images['Seg']),
+                                                    np.array(images['Marker'].convert('L')), resolution,
+                                                    seg_thresh, size_thresh, marker_thresh, size_thresh_upper)
+        processed_images = {}
+        processed_images['SegOverlaid'] = Image.fromarray(overlay)
+        processed_images['SegRefined'] = Image.fromarray(refined)
+        return processed_images, scoring
+
+    elif opt.model == 'DeepLIIFExt':
+        resolution = '40x' if tile_size > 768 else ('20x' if tile_size > 384 else '10x')
+        processed_images = {}
+        scoring = {}
+        for img_name in list(images.keys()):
+            if 'Seg' in img_name:
+                seg_img = images[img_name]
+                overlay, refined, score = compute_results(np.array(orig), np.array(images[img_name]),
+                                                          None, resolution,
+                                                          seg_thresh, size_thresh, marker_thresh, size_thresh_upper)
+    
+                processed_images[img_name + '_Overlaid'] = Image.fromarray(overlay)
+                processed_images[img_name + '_Refined'] = Image.fromarray(refined)
+                scoring[img_name] = score
+        return processed_images, scoring
 
-    return images, scoring
+    else:
+        raise Exception(f'postprocess() not implemented for model {opt.model}')
 
 
 def infer_modalities(img, tile_size, model_dir, eager_mode=False,
-                     color_dapi=False, color_marker=False):
+                     color_dapi=False, color_marker=False, opt=None):
     """
     This function is used to infer modalities for the given image using a trained model.
     :param img: The input image.
@@ -379,6 +464,11 @@ def infer_modalities(img, tile_size, model_dir, eager_mode=False,
     :param model_dir: The directory containing serialized model files.
     :return: The inferred modalities and the segmentation mask.
     """
+    if opt is None:
+        opt = get_opt(model_dir)
+        opt.use_dp = False
+        print_options(opt)
+    
     if not tile_size:
         tile_size = check_multi_scale(Image.open('./images/target.png').convert('L'),
                                       img.convert('L'))
@@ -391,12 +481,16 @@ def infer_modalities(img, tile_size, model_dir, eager_mode=False,
         model_path=model_dir,
         eager_mode=eager_mode,
         color_dapi=color_dapi,
-        color_marker=color_marker
+        color_marker=color_marker,
+        opt=opt
     )
-
-    post_images, scoring = postprocess(img, images['Seg'], small_object_size=20)
-    images = {**images, **post_images}
-    return images, scoring
+    
+    if not hasattr(opt,'seg_gen') or (hasattr(opt,'seg_gen') and opt.seg_gen): # the first condition accounts for old settings of deepliif; the second refers to deepliifext models
+        post_images, scoring = postprocess(img, images, tile_size, opt=opt)
+        images = {**images, **post_images}
+        return images, scoring
+    else:
+        return images, None
 
 
 def infer_results_for_wsi(input_dir, filename, output_dir, model_dir, tile_size, region_size=20000):

deepliif/models/base_model.py

@@ -4,6 +4,8 @@ from collections import OrderedDict
 from abc import ABC, abstractmethod
 from . import networks
 from ..util import disable_batchnorm_tracking_stats
+from deepliif.util import *
+import itertools
 
 
 class BaseModel(ABC):
@@ -35,8 +37,9 @@ class BaseModel(ABC):
         self.is_train = opt.is_train
         self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')  # get device name: CPU or GPU
         self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)  # save all the checkpoints to save_dir
-        if opt.preprocess != 'scale_width':  # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
+        if opt.phase == 'train' and opt.preprocess != 'scale_width': # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
             torch.backends.cudnn.benchmark = True
+        # especially for inference, cudnn benchmark can cause excessive usage of GPU memory for the first image in the sequence in order to find the best conv alg which is not necessary.
         self.loss_names = []
         self.model_names = []
         self.visual_names = []
@@ -89,7 +92,10 @@ class BaseModel(ABC):
         """Make models eval mode during test time"""
         for name in self.model_names:
             if isinstance(name, str):
-                net = getattr(self, 'net' + name)
+                if '_' in name:
+                    net = getattr(self, 'net' + name.split('_')[0])[int(name.split('_')[-1]) - 1]
+                else:
+                    net = getattr(self, 'net' + name)
                 net.eval()
                 net = disable_batchnorm_tracking_stats(net)
 
@@ -127,7 +133,13 @@ class BaseModel(ABC):
         visual_ret = OrderedDict()
         for name in self.visual_names:
             if isinstance(name, str):
-                visual_ret[name] = getattr(self, name)
+                if not hasattr(self, name):
+                    if len(name.split('_')) == 2:
+                        visual_ret[name] = getattr(self, name.split('_')[0])[int(name.split('_')[-1]) -1]
+                    else:
+                        visual_ret[name] = getattr(self, name.split('_')[0] + '_' + name.split('_')[1])[int(name.split('_')[-1]) - 1]
+                else:
+                    visual_ret[name] = getattr(self, name)
         return visual_ret
 
     def get_current_losses(self):
@@ -135,7 +147,16 @@ class BaseModel(ABC):
         errors_ret = OrderedDict()
         for name in self.loss_names:
             if isinstance(name, str):
-                errors_ret[name] = float(getattr(self, 'loss_' + name))  # float(...) works for both scalar tensor and float number
+                if not hasattr(self, 'loss_'+name): # appears in DeepLIIFExt
+                    if len(name.split('_')) == 2:
+                        errors_ret[name] = float(getattr(self, 'loss_' + name.split('_')[0])[int(
+                            name.split('_')[-1]) - 1])  # float(...) works for both scalar tensor and float number
+                    else:
+                        errors_ret[name] = float(getattr(self, 'loss_' + name.split('_')[0] + '_' + name.split('_')[1])[int(
+                            name.split('_')[-1]) - 1])  # float(...) works for both scalar tensor and float number
+                else: # single numeric value
+                    errors_ret[name] = float(getattr(self, 'loss_' + name)) 
+                    
         return errors_ret
 
     def save_networks(self, epoch, save_from_one_process=False):
@@ -151,7 +172,10 @@ class BaseModel(ABC):
             if isinstance(name, str):
                 save_filename = '%s_net_%s.pth' % (epoch, name)
                 save_path = os.path.join(self.save_dir, save_filename)
-                net = getattr(self, 'net' + name)
+                if '_' in name:
+                    net = getattr(self, 'net' + name.split('_')[0])[int(name.split('_')[-1]) - 1]
+                else:
+                    net = getattr(self, 'net' + name)
 
                 if len(self.gpu_ids) > 0 and torch.cuda.is_available():
                     torch.save(net.module.cpu().state_dict(), save_path)
@@ -219,13 +243,32 @@ class BaseModel(ABC):
             if isinstance(name, str):
                 load_filename = '%s_net_%s.pth' % (epoch, name)
                 load_path = os.path.join(self.save_dir, load_filename)
-                net = getattr(self, 'net' + name)
+                if '_' in name:
+                    net = getattr(self, 'net' + name.split('_')[0])[int(name.split('_')[-1]) - 1]
+                else:
+                    net = getattr(self, 'net' + name)
                 if isinstance(net, torch.nn.DataParallel):
                     net = net.module
+                
+                self.set_requires_grad(net,self.opt.is_train)
+                # check if gradients are disabled
+                names_layer_requires_grad = []
+                for name, param in net.named_parameters():
+                    if param.requires_grad:
+                        names_layer_requires_grad.append(name)
+                
                 print('loading the model from %s' % load_path)
                 # if you are using PyTorch newer than 0.4 (e.g., built from
                 # GitHub source), you can remove str() on self.device
-                state_dict = torch.load(load_path, map_location=str(self.device))
+                
+                if self.opt.is_train or self.opt.use_dp:
+                    device = self.device
+                else:
+                    device = torch.device('cpu') # load in cpu first; later in __inite__.py::init_nets we will move it to the specified device
+                    
+                net.to(device)
+                state_dict = torch.load(load_path, map_location=str(device))
+                
                 if hasattr(state_dict, '_metadata'):
                     del state_dict._metadata
 
@@ -243,7 +286,10 @@ class BaseModel(ABC):
         print('---------- Networks initialized -------------')
         for name in self.model_names:
             if isinstance(name, str):
-                net = getattr(self, 'net' + name)
+                if '_' in name:
+                    net = getattr(self, 'net' + name.split('_')[0])[int(name.split('_')[-1]) - 1]
+                else:
+                    net = getattr(self, 'net' + name)
                 num_params = 0
                 for param in net.parameters():
                     num_params += param.numel()