Nonsensical Unet output with model.eval()

username: ElPolloDiablo
cooked: <p>Hi Forum,</p>
<p>I’m having an issue with a pretrained model (Unet, librabry: <a href="https://github.com/qubvel/segmentation_models.pytorch" rel="noopener nofollow ugc">Pytorch segmentation models</a>) which I use for the tile-wise segmentation of large (6000x6000) histological images. The model performs well during training and testing, whereas I use model.train() for the training step and model.eval() for the testing step. No unexpected behavior here. I use a batch size of 20 as my GPU only has 8GB RAM. (For snippet of training code, see below)</p>
<p>In a second step, I apply the model to a large amount of images in a separate script. Now, after loading the model with the following code, I get purely nonsensical results (see images below) - when I omit <code>model.eval()</code>, the result looks good. I still get some batch effects, which I can mitigate a bit with proper shuffling.</p>
<pre><code class="lang-auto">    model = smp.Unet(
        encoder_name='resnet50', 
        encoder_weights='imagenet', 
        classes=3, 
        activation=None,
    )
    
    model.load_state_dict(torch.load(BST_MODEL))
    model = model.to(DEVICE)
    model.eval()
</code></pre>
<p>Here’s a collage of the (histological) images I’m working on - I am aware that it probably won’t mean much to you - but I hope it illustrates the problem better than loss metrics.<br>
|<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/7/67bc5e4cd70aeea8f6be9bbd9152b5087199a63f.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/67bc5e4cd70aeea8f6be9bbd9152b5087199a63f" title="N182a_SAS_1_1_170905-Scene-1-ScanRegion0_HE.czi - N182a_SAS_1_1_170905-Scene-1-ScanRegion0_HE.czi #3-1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/7/67bc5e4cd70aeea8f6be9bbd9152b5087199a63f_2_225x225.jpeg" alt="N182a_SAS_1_1_170905-Scene-1-ScanRegion0_HE.czi - N182a_SAS_1_1_170905-Scene-1-ScanRegion0_HE.czi #3-1" data-base62-sha1="eNGKVa7klwyelbfBN2OLWpS8HeL" width="225" height="225" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/7/67bc5e4cd70aeea8f6be9bbd9152b5087199a63f_2_225x225.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/7/67bc5e4cd70aeea8f6be9bbd9152b5087199a63f_2_337x337.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/7/67bc5e4cd70aeea8f6be9bbd9152b5087199a63f_2_450x450.jpeg 2x" data-dominant-color="958298"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">N182a_SAS_1_1_170905-Scene-1-ScanRegion0_HE.czi - N182a_SAS_1_1_170905-Scene-1-ScanRegion0_HE.czi #3-1</span><span class="informations">3590×2884 2.27 MB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div>  <div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/8/183459a3f78e28e6f197968f721c29fca85880dd.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/183459a3f78e28e6f197968f721c29fca85880dd" title="HE_seg_DL_modelTRAIN"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/8/183459a3f78e28e6f197968f721c29fca85880dd_2_225x225.jpeg" alt="HE_seg_DL_modelTRAIN" data-base62-sha1="3s7Cno26gkikbClHwVoSVSCiugR" width="225" height="225" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/8/183459a3f78e28e6f197968f721c29fca85880dd_2_225x225.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/1/8/183459a3f78e28e6f197968f721c29fca85880dd_2_337x337.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/8/183459a3f78e28e6f197968f721c29fca85880dd_2_450x450.jpeg 2x" data-dominant-color="4C8759"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">HE_seg_DL_modelTRAIN</span><span class="informations">3590×2884 1.19 MB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div> <div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/9/d9069d14916e86cbd52eb8f7d9dcf723ad43eda8.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/d9069d14916e86cbd52eb8f7d9dcf723ad43eda8" title="HE_seg_DL_modelEVAL"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d9069d14916e86cbd52eb8f7d9dcf723ad43eda8_2_225x225.jpeg" alt="HE_seg_DL_modelEVAL" data-base62-sha1="uXTMoOEyBlGArKl09XZP3NfglEk" width="225" height="225" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d9069d14916e86cbd52eb8f7d9dcf723ad43eda8_2_225x225.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d9069d14916e86cbd52eb8f7d9dcf723ad43eda8_2_337x337.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d9069d14916e86cbd52eb8f7d9dcf723ad43eda8_2_450x450.jpeg 2x" data-dominant-color="CDDCD3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">HE_seg_DL_modelEVAL</span><span class="informations">3590×2884 348 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The difference in code leading to the two outputs is exclusively the .eval() command after moving the model to the GPU.</p>
<p>I suspect that the issue is closely related to things that have been discussed in other threads (like <a href="https://discuss.pytorch.org/t/model-eval-gives-incorrect-loss-for-model-with-batchnorm-layers/7561">this</a> or <a href="https://discuss.pytorch.org/t/what-does-model-eval-do-for-batchnorm-layer/7146">this</a>).</p>
<p>Things I have tried so far:</p>
<ul>
<li>track_running_stats=False for batchnorm2d layers</li>
<li>varying the momentum parameter for batchnorm2d layers</li>
<li>Increasing batch size to 32</li>
</ul>
<p>Note: Since the images are too large, I process the whole images as 256x256-sized tiles, which I feed as (shuffled) batches through the model. I process the whole image several times with offset tile coordinates. The multiple predictions are then averaged to suppress tile-artifacts.</p>
<p>My code is completely available <a href="https://github.com/jo-mueller/NecroticSegmentation" rel="noopener nofollow ugc">here</a>. As stated above, train/testing (Train.py) works fine, things go sideways after loading and running the model (Process.py).</p>
<p>I guess I am doing some rookie mistake but I cannot identify the error.<br>
<strong>Thank you</strong> in advance for your time and help!</p>
<p>PS: Train/Test code snippet:</p>
<pre><code class="lang-auto">    # dataloaders &amp; PerformanceMeter
    train_dataloader = DataLoader(train_dataset, BATCH_SIZE, shuffle=True, num_workers=num_workers)
    val_dataloader   = DataLoader(val_dataset, BATCH_SIZE, shuffle=True, num_workers=num_workers)
    Monitor = helper_functions.PerformanceMeter()
    
    # score lists
    train_score = []
    test_score = []
    
    # train
    for epoch in range(EPOCHS):
        
        model.train()
        optimizer.zero_grad()
        tk0 = tqdm(train_dataloader, total=len(train_dataloader))
        
        for b_idx, data in enumerate(tk0):
            
            # move images on GPU
            for key, value in data.items():
                data[key] = value.to(device).float()
                
            # train
            data['prediction']  = model(data['image'])
            loss = criterion(data['prediction'], data['mask'][:, 0].long())
            loss.backward()
            optimizer.step()
            tk0.set_postfix(loss=loss.cpu().detach().numpy())

        # evaluate
        model.eval()
        with torch.no_grad():
            
            tk1 = tqdm(val_dataloader, total=len(val_dataloader))
            dice = np.zeros(N_CLASSES)
            for b_idx, data in enumerate(tk1):
                
                # Eval
                data['prediction'] = model(data['image'].to(device).float())                 
                out = torch.argmax(data['prediction'], dim=1).view(-1)
                mask = data['mask'].view(-1)
                
                score = jaccard_score(mask.cpu(), out.cpu(), average=None)
                dice += score

</code></pre>