Different predictions from a trained model every time a Jupyter kernel is restarted
🐛 Bug
This pertains to using the forward pass for making inference. For a given input image and fixed model, the output values change for both train and eval modes. Note: The model here is fixed and not being trained i.e. I'm just doing inference from a checkpoint.
The lightning module for MNIST classification:
from torchvision.datasets import FashionMNIST, MNIST
import matplotlib.pyplot as plt
import torch
import numpy as np
import os
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
import pytorch_lightning as pl
from pytorch_lightning.metrics.functional import accuracy
%matplotlib inline
class ANN(pl.LightningModule):
    def __init__(self, data_dir='./', learning_rate=2e-4):

        super().__init__()

        # Set our init args as class attributes
        self.data_dir = data_dir
        self.learning_rate = learning_rate
        self.test_targets = []
        self.test_embeddings = torch.zeros((0, 100),
                                           dtype=torch.float32,
                                           device='cuda:0')
        self.test_predictions = []
        # Hardcode some dataset specific attributes
        self.num_classes = 10
        self.dims = (1, 28, 28)
        self.transform = transforms.Compose([transforms.ToTensor()])
        self.conv1 = nn.Conv2d(1, 16, 3)
        self.bn1 = nn.BatchNorm2d(16)
        self.maxpool1 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(16 * 13 * 13, 100)
        self.fc2 = nn.Linear(100, self.num_classes)
        # Define PyTorch model

    def forward(self, x):
        #         print(x.shape)
        x = F.relu(self.bn1(self.conv1(x)))
        #         print(x.shape)
        x = F.relu(self.maxpool1(x))
        # x = F.relu(self.bn2(self.conv2(x)))
        #         print(x.shape)
        x = x.view(-1, 16 * 13 * 13)
        x = self.fc1(x)
        y = self.fc2(F.relu(x))
        return x, y

    def training_step(self, batch, batch_idx):
        x, y = batch
        _, logits = self(x)
        loss = F.cross_entropy(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = accuracy(preds, y)
        self.log('train_loss', loss, prog_bar=True)
        self.log('train_acc', acc, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        _, logits = self(x)
        loss = F.cross_entropy(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = accuracy(preds, y)

        # Calling self.log will surface up scalars for you in TensorBoard
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', acc, prog_bar=True)
        return loss

    def test_step(self, batch, batch_idx):
        # Here we just reuse the validation_step for testing
        x, y = batch
        embeddings, logits = self(x)
        loss = F.cross_entropy(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = accuracy(preds, y)
        self.test_predictions.extend(preds.detach().cpu().tolist())
        self.test_targets.extend(y.detach().cpu().tolist())
        self.test_embeddings = torch.cat((self.test_embeddings, embeddings), 0)
        self.log('test_acc', acc)
        self.log('test_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters())
        return optimizer

    ####################
    # DATA RELATED HOOKS
    ####################

    def prepare_data(self):
        # download
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):

        # Assign train/val datasets for use in dataloaders
        if stage == 'fit':
            dataset_full = MNIST(self.data_dir,
                                 train=True,
                                 transform=self.transform)
            self.dataset_train, self.dataset_val = random_split(
                dataset_full, [55000, 5000])

        # Assign test dataset for use in dataloader(s)
        if stage == 'test':
            self.dataset_test = MNIST(self.data_dir,
                                      train=False,
                                      transform=self.transform)
    def train_dataloader(self):
        return DataLoader(self.dataset_train, batch_size=32, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.dataset_val, batch_size=32)

    def test_dataloader(self):
        return DataLoader(self.dataset_test, batch_size=32)
model = ANN()

### To Reproduce
model.load_from_checkpoint(r".\lightning_logs\version_5\checkpoints\epoch=4-step=8594.ckpt")
model = model.cuda()

transform = transforms.Compose([transforms.ToTensor()])
data = MNIST(".", train=False, download=True, transform=transform)
test_sample = torch.unsqueeze(data.data[:32].clone().cuda(),1).float()
plt.figure()
for i in range(32):
    plt.subplot(4,8,i+1)
    plt.imshow(data.data[i])
    plt.axis('off')
out = model(test_sample)
print(out[1][0])

with torch.no_grad():
    out = model(test_sample)
print(out[1][0])

model.eval()

Train behaviour:

Eval behaviour:

Expected behavior
The predictions should  be same every time I do a forward pass. I understand the train and eval behaviour is different. However, given a mode (train/eval) for a fixed input the model should give the same output every time. I spent almost an entire day trying to find out what is going wrong. This is a very troublesome behaviour. I've not encountered this in PyTorch before.
Environment
name: PyTorch
channels:
  - pytorch
  - conda-forge
  - defaults
dependencies:
  - absl-py=0.11.0=py37h03978a9_0
  - aiohttp=3.7.3=py37hcc03f2d_1
  - alabaster=0.7.12=py_0
  - appdirs=1.4.4=pyh9f0ad1d_0
  - argh=0.26.2=pyh9f0ad1d_1002
  - argon2-cffi=20.1.0=py37hcc03f2d_2
  - asciimatics=1.12.0=py37h03978a9_0
  - astroid=2.4.2=py37hc8dfbb8_1
  - async-timeout=3.0.1=py_1000
  - async_generator=1.10=py_0
  - atomicwrites=1.4.0=pyh9f0ad1d_0
  - attrs=20.3.0=pyhd3deb0d_0
  - autopep8=1.5.4=pyh9f0ad1d_0
  - babel=2.9.0=pyhd3deb0d_0
  - backcall=0.2.0=pyh9f0ad1d_0
  - backports=1.0=py_2
  - backports.functools_lru_cache=1.6.1=py_0
  - bcrypt=3.2.0=py37hcc03f2d_1
  - beautifulsoup4=4.9.3=pyhb0f4dca_0
  - black=20.8b1=py_1
  - blas=2.107=mkl
  - blas-devel=3.9.0=7_mkl
  - bleach=3.2.1=pyh9f0ad1d_0
  - blinker=1.4=py_1
  - blosc=1.21.0=h0e60522_0
  - boto3=1.16.50=pyhd8ed1ab_0
  - botocore=1.19.53=pyhd8ed1ab_0
  - brotli=1.0.9=h0e60522_4
  - brotlipy=0.7.0=py37h0013d47_1001
  - bzip2=1.0.8=h8ffe710_4
  - ca-certificates=2020.12.5=h5b45459_0
  - cachetools=4.1.1=py_0
  - charls=2.2.0=h39d44d4_0
  - click=7.1.2=pyh9f0ad1d_0
  - colorama=0.4.4=pyh9f0ad1d_0
  - configobj=5.0.6=py_0
  - configparser=5.0.1=py_0
  - contextlib2=0.5.5=py_2
  - cryptography=3.3.1=py37hd8e9650_1
  - cudatoolkit=10.1.243=h3826478_6
  - cycler=0.10.0=py_2
  - cytoolz=0.11.0=py37hcc03f2d_1
  - dask-core=2020.12.0=pyhd8ed1ab_0
  - dataclasses=0.7=pyhb2cacf7_7
  - decorator=4.4.2=py_0
  - defusedxml=0.6.0=py_0
  - diff-match-patch=20200713=pyh9f0ad1d_0
  - distro=1.5.0=pyh9f0ad1d_0
  - docutils=0.16=py37h03978a9_3
  - dvc=0.54.1=py37_0
  - entrypoints=0.3=pyhd8ed1ab_1003
  - et_xmlfile=1.0.1=py_1001
  - flake8=3.8.4=py_0
  - freeglut=3.2.1=h0e60522_0
  - freetype=2.10.4=h546665d_0
  - fsspec=0.8.5=pyhd8ed1ab_0
  - funcy=1.15=pyh9f0ad1d_0
  - future=0.18.2=py37h03978a9_3
  - geos=3.8.1=he025d50_0
  - giflib=5.2.1=h8d14728_2
  - gitdb=4.0.5=py_0
  - google-auth=1.24.0=pyhd3deb0d_0
  - google-auth-oauthlib=0.4.1=py_2
  - grandalf=0.6=py_0
  - grpcio=1.34.0=py37h04d2302_1
  - helpdev=0.7.1=pyhd8ed1ab_0
  - humanize=3.2.0=pyhd8ed1ab_0
  - icu=68.1=h0e60522_0
  - idna=2.10=pyh9f0ad1d_0
  - ignite=0.3.0=py37_0
  - imagecodecs=2021.1.11=py37hfc4c502_1
  - imageio=2.9.0=py_0
  - imagesize=1.2.0=py_0
  - imgaug=0.4.0=py_1
  - importlib-metadata=3.4.0=py37h03978a9_0
  - importlib_metadata=3.4.0=hd8ed1ab_0
  - inflect=5.0.2=py37h03978a9_1
  - iniconfig=1.1.1=pyh9f0ad1d_0
  - intel-openmp=2020.3=h57928b3_311
  - intervaltree=3.0.2=py_0
  - ipykernel=5.4.2=py37h7b7c402_0
  - ipympl=0.6.2=pyhd8ed1ab_0
  - ipython=7.19.0=py37heaed05f_1
  - ipython_genutils=0.2.0=py_1
  - ipywidgets=7.6.3=pyhd3deb0d_0
  - isort=5.7.0=pyhd8ed1ab_0
  - jasper=2.0.14=h77af90b_2
  - jdcal=1.4.1=py_0
  - jedi=0.14.1=py37_0
  - jinja2=2.11.2=pyh9f0ad1d_0
  - jmespath=0.10.0=pyh9f0ad1d_0
  - joblib=1.0.0=pyhd8ed1ab_0
  - jpeg=9d=h8ffe710_0
  - jsonpath-ng=1.5.2=pyh9f0ad1d_0
  - jsonschema=3.2.0=py_2
  - jupyter_client=6.1.11=pyhd8ed1ab_1
  - jupyter_core=4.7.0=py37h03978a9_0
  - jupyterlab_pygments=0.1.2=pyh9f0ad1d_0
  - jupyterlab_widgets=1.0.0=pyhd8ed1ab_1
  - jxrlib=1.1=h8ffe710_2
  - keyring=21.8.0=py37h03978a9_0
  - kiwisolver=1.3.1=py37h8c56517_1
  - lazy-object-proxy=1.4.3=py37h4ab8f01_2
  - lcms2=2.11=h6a326e1_1
  - lerc=2.2.1=h0e60522_0
  - libaec=1.0.4=h39d44d4_1
  - libblas=3.9.0=7_mkl
  - libcblas=3.9.0=7_mkl
  - libclang=11.0.1=default_h5c34c98_0
  - libdeflate=1.7=h8ffe710_4
  - libiconv=1.16=he774522_0
  - liblapack=3.9.0=7_mkl
  - liblapacke=3.9.0=7_mkl
  - libopencv=4.5.0=py37ha0199f4_7
  - libpng=1.6.37=h1d00b33_2
  - libprotobuf=3.14.0=h7755175_0
  - libsodium=1.0.18=h8d14728_1
  - libspatialindex=1.9.3=he025d50_3
  - libtiff=4.2.0=hc10be44_0
  - libuv=1.40.0=h8ffe710_0
  - libwebp-base=1.1.0=h8ffe710_3
  - libxml2=2.9.10=hf5bbc77_3
  - libxslt=1.1.33=h65864e5_2
  - libzopfli=1.0.3=h0e60522_0
  - lxml=4.6.2=py37hd07aab1_0
  - lz4-c=1.9.3=h8ffe710_0
  - m2w64-gcc-libgfortran=5.3.0=6
  - m2w64-gcc-libs=5.3.0=7
  - m2w64-gcc-libs-core=5.3.0=7
  - m2w64-gmp=6.1.0=2
  - m2w64-libwinpthread-git=5.0.0.4634.697f757=2
  - markdown=3.3.3=pyh9f0ad1d_0
  - markupsafe=1.1.1=py37hcc03f2d_3
  - matplotlib-base=3.3.3=py37h3379fd5_0
  - mccabe=0.6.1=py_1
  - mistune=0.8.4=py37hcc03f2d_1003
  - mkl=2020.4=hb70f87d_311
  - mkl-devel=2020.4=h57928b3_312
  - mkl-include=2020.4=hb70f87d_311
  - more-itertools=8.6.0=pyhd8ed1ab_0
  - msys2-conda-epoch=20160418=1
  - multidict=5.1.0=py37hcc03f2d_1
  - mypy_extensions=0.4.3=py37h03978a9_3
  - nanotime=0.5.2=py_0
  - nbclient=0.5.1=py_0
  - nbconvert=6.0.7=py37h03978a9_3
  - nbformat=5.0.8=py_0
  - nest-asyncio=1.4.3=pyhd8ed1ab_0
  - networkx=2.5=py_0
  - ninja=1.10.2=h5362a0b_0
  - nose=1.3.7=py_1006
  - notebook=6.1.6=py37h03978a9_0
  - numpy=1.19.5=py37hd20adf4_1
  - numpydoc=1.1.0=py_1
  - olefile=0.46=pyh9f0ad1d_1
  - opencv=4.5.0=py37h03978a9_7
  - openjpeg=2.4.0=h48faf41_0
  - openpyxl=3.0.5=py_0
  - openssl=1.1.1i=h8ffe710_0
  - packaging=20.8=pyhd3deb0d_0
  - pandas=1.2.0=py37h08fd248_0
  - pandoc=2.11.3.2=h8ffe710_0
  - pandocfilters=1.4.2=py_1
  - paramiko=2.7.2=pyh9f0ad1d_0
  - parso=0.5.2=py_0
  - pathspec=0.8.1=pyhd3deb0d_0
  - pexpect=4.8.0=pyh9f0ad1d_2
  - pickleshare=0.7.5=py_1003
  - pillow=8.1.0=py37hdec93a9_1
  - pip=20.3.3=pyhd8ed1ab_0
  - pluggy=0.13.1=py37h03978a9_4
  - ply=3.11=py_1
  - pooch=1.3.0=pyhd8ed1ab_0
  - prometheus_client=0.9.0=pyhd3deb0d_0
  - prompt-toolkit=3.0.10=pyha770c72_0
  - protobuf=3.14.0=py37hf2a7229_1
  - psutil=5.8.0=py37hcc03f2d_1
  - ptyprocess=0.7.0=pyhd3deb0d_0
  - py=1.10.0=pyhd3deb0d_0
  - py-opencv=4.5.0=py37heaed05f_7
  - pyasn1=0.4.8=py_0
  - pyasn1-modules=0.2.7=py_0
  - pycodestyle=2.6.0=pyh9f0ad1d_0
  - pycparser=2.20=pyh9f0ad1d_2
  - pydocstyle=5.1.1=py_0
  - pyfiglet=0.8.post1=py_0
  - pyflakes=2.2.0=pyh9f0ad1d_0
  - pygments=2.7.3=pyhd8ed1ab_0
  - pyjwt=2.0.0=pyhd8ed1ab_0
  - pylint=2.6.0=py37hc8dfbb8_1
  - pynacl=1.4.0=py37h230ea1a_2
  - pyodbc=4.0.30=py37h1834ac0_1
  - pyopenssl=20.0.1=pyhd8ed1ab_0
  - pyparsing=2.4.7=pyh9f0ad1d_0
  - pyqt=5.12.3=py37h03978a9_7
  - pyqt-impl=5.12.3=py37hf2a7229_7
  - pyqt5-sip=4.19.18=py37hf2a7229_7
  - pyqtchart=5.12=py37hf2a7229_7
  - pyqtwebengine=5.12.1=py37hf2a7229_7
  - pyrsistent=0.17.3=py37hcc03f2d_2
  - pysocks=1.7.1=py37h03978a9_3
  - pytest=6.2.1=py37h03978a9_0
  - python=3.7.9=h7840368_0_cpython
  - python-dateutil=2.8.1=py_0
  - python-jsonrpc-server=0.3.4=pyh9f0ad1d_1
  - python-language-server=0.31.7=py37_0
  - python_abi=3.7=1_cp37m
  - pytorch=1.7.0=py3.7_cuda101_cudnn7_0
  - pytorch-lightning=1.1.5=pyhd8ed1ab_0
  - pytz=2020.5=pyhd8ed1ab_0
  - pywavelets=1.1.1=py37hda49f71_3
  - pywin32=228=py37hcc03f2d_1
  - pywin32-ctypes=0.2.0=py37h03978a9_1003
  - pywinpty=0.5.7=py37hc8dfbb8_1
  - pyzmq=20.0.0=py37h0d95fc2_1
  - qdarkstyle=2.8.1=pyhd8ed1ab_2
  - qt=5.12.9=h5909a2a_3
  - qtawesome=1.0.2=pyhd8ed1ab_0
  - qtconsole=5.0.1=pyhd8ed1ab_0
  - qtpy=1.9.0=py_0
  - regex=2020.11.13=py37hcc03f2d_1
  - requests=2.25.1=pyhd3deb0d_0
  - requests-oauthlib=1.3.0=pyh9f0ad1d_0
  - rope=0.18.0=pyh9f0ad1d_0
  - rsa=4.7=pyhd3deb0d_0
  - rtree=0.9.7=py37h13cc57e_1
  - ruamel=1.0=py37h03978a9_4
  - ruamel.yaml=0.15.100=py37hfa6e2cd_0
  - s3transfer=0.3.4=pyhd8ed1ab_0
  - schema=0.7.3=pyh9f0ad1d_0
  - scikit-image=0.18.1=py37h08fd248_0
  - scikit-learn=0.24.0=py37heb15398_0
  - scipy=1.6.0=py37h6db1a17_0
  - send2trash=1.5.0=py_0
  - setuptools=49.6.0=py37h03978a9_3
  - shapely=1.7.1=py37h82f916b_1
  - shortuuid=1.0.1=py37h03978a9_4
  - six=1.15.0=pyh9f0ad1d_0
  - smmap=3.0.4=pyh9f0ad1d_0
  - snappy=1.1.8=ha925a31_3
  - snowballstemmer=2.0.0=py_0
  - sortedcontainers=2.3.0=pyhd8ed1ab_0
  - soupsieve=2.0.1=py_1
  - sphinx=3.4.3=pyhd8ed1ab_0
  - sphinxcontrib-applehelp=1.0.2=py_0
  - sphinxcontrib-devhelp=1.0.2=py_0
  - sphinxcontrib-htmlhelp=1.0.3=py_0
  - sphinxcontrib-jsmath=1.0.1=py_0
  - sphinxcontrib-qthelp=1.0.3=py_0
  - sphinxcontrib-serializinghtml=1.1.4=py_0
  - spyder=4.0.1=py37_0
  - spyder-kernels=1.8.1=py37_0
  - sqlite=3.34.0=h8ffe710_0
  - tensorboard=2.4.0=pyhd8ed1ab_0
  - tensorboard-plugin-wit=1.7.0=pyh9f0ad1d_0
  - terminado=0.9.2=py37h03978a9_0
  - testpath=0.4.4=py_0
  - threadpoolctl=2.1.0=pyh5ca1d4c_0
  - tk=8.6.10=h8ffe710_1
  - toml=0.10.2=pyhd8ed1ab_0
  - toolz=0.11.1=py_0
  - torchaudio=0.7.0=py37
  - torchtext=0.6.0=py_1
  - torchvision=0.8.1=py37_cu101
  - tornado=6.1=py37hcc03f2d_1
  - tqdm=4.55.1=pyhd8ed1ab_0
  - traitlets=5.0.5=py_0
  - treelib=1.6.1=pyh9f0ad1d_0
  - typed-ast=1.4.2=py37hcc03f2d_0
  - typing-extensions=3.7.4.3=0
  - typing_extensions=3.7.4.3=py_0
  - ujson=1.35=py37hf2a7229_1003
  - urllib3=1.26.2=pyhd8ed1ab_0
  - vc=14.2=hb210afc_2
  - vs2015_runtime=14.28.29325=h5e1d092_0
  - watchdog=1.0.2=py37h03978a9_1
  - wcwidth=0.2.5=pyh9f0ad1d_2
  - webencodings=0.5.1=py_1
  - werkzeug=1.0.1=pyh9f0ad1d_0
  - wheel=0.36.2=pyhd3deb0d_0
  - widgetsnbextension=3.5.1=py37h03978a9_4
  - win-unicode-console=0.5=py37_0
  - win_inet_pton=1.1.0=py37h03978a9_2
  - wincertstore=0.2=py37h03978a9_1006
  - winpty=0.4.3=4
  - wrapt=1.11.2=py37h4ab8f01_1
  - xlrd=1.2.0=pyh9f0ad1d_1
  - xz=5.2.5=h62dcd97_1
  - yaml=0.2.5=he774522_0
  - yapf=0.30.0=pyh9f0ad1d_0
  - yarl=1.6.3=py37hcc03f2d_1
  - zc.lockfile=2.0=py_0
  - zeromq=4.3.3=h0e60522_3
  - zfp=0.5.5=h0e60522_4
  - zipp=3.4.0=py_0
  - zlib=1.2.11=h62dcd97_1010
  - zstd=1.4.8=h4e2f164_1
  - pip:
    - alembic==1.4.2
    - certifi==2020.12.5
    - cffi==1.14.4
    - chardet==4.0.0
    - cloudpickle==1.3.0
    - cython==0.29.21
    - databricks-cli==0.11.0
    - docker==4.2.1
    - flask==1.1.2
    - gitpython==3.1.3
    - gorilla==0.3.0
    - gym==0.17.2
    - isodate==0.6.0
    - itsdangerous==1.1.0
    - mako==1.1.3
    - mlflow==1.8.0
    - oauthlib==3.1.0
    - opencv-python==4.5.1.48
    - prometheus-flask-exporter==0.13.0
    - pycocotools==2.0.2
    - pyglet==1.5.0
    - pypiwin32==223
    - python-editor==1.0.4
    - pyyaml==5.3.1
    - querystring-parser==1.2.4
    - simplejson==3.17.0
    - sqlalchemy==1.3.13
    - sqlparse==0.3.1
    - tabulate==0.8.7
    - tifffile==2020.6.3
    - waitress==1.4.4
    - websocket-client==0.57.0
prefix: C:\Users\manpreet\Anaconda3\envs\PyTorch


PyTorch Version (e.g., 1.0): 1.7.0
OS (e.g., Linux): Windows 10
How you installed PyTorch (conda, pip, source): conda
Build command you used (if compiling from source): NA
Python version: 3.7.9
CUDA/cuDNN version: 10.1