add tomography splitting loss example

docs/search.json

@@ -6,6 +6,20 @@
     "section": "",
     "text": "%cd ..\n\nc:\\Users\\s2558406\\Documents\\Repos\\cmr-experiments\n\n\nc:\\Users\\s2558406\\Documents\\Repos\\cmr-experiments\\venv\\Lib\\site-packages\\IPython\\core\\magics\\osm.py:417: UserWarning: This is now an optional IPython functionality, setting dhist requires you to install the `pickleshare` library.\n  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n\n\n\nfrom utils.edipo.data.transforms import CineNetDataTransform\nfrom utils.edipo.data.mri_data import RawDataSample\nfrom utils.edipo.models.crnn import CRNN\nfrom utils import *\nfrom pathlib import Path\nfrom torch.utils.data import DataLoader, random_split\nimport torch\nimport numpy as np\n\n\n# ONLY RUN when new dataset_cache run on Linux\nimport pathlib, pickle\nwith open(\"dataset_cache.pkl\", \"rb\") as f:\n    with set_posix_windows():\n        dataset_cache = pickle.load(f)\n    \n    new_cache = {\n        pathlib.WindowsPath(r\"M:\\data\\CMRxRecon\\SingleCoil\\Cine\\TrainingSet\"):\n    [\n        RawDataSample(r.fname.replace(\"/home/s2558406/RDS\", \"M:\").replace(\"/\", \"\\\\\"), r.slice_ind, r.metadata) for r in dataset_cache[\n            pathlib.WindowsPath('/home/s2558406/RDS/data/CMRxRecon/SingleCoil/Cine/TrainingSet')\n        ]\n    ]}\n\nwith open(\"dataset_cache_windows.pkl\", \"wb\") as f:\n    pickle.dump(new_cache, f)\n    \n\n\ndataset = DeepinvSliceDataset(\n    Path(r\"M:\\data\\CMRxRecon\"), \n    transform=CineNetDataTransform(time_window=12, apply_mask=True, normalize=False), \n    set_name=\"TrainingSet\",\n    acc_folders=[\"FullSample\"],\n    mask_folder=\"TimeVaryingGaussianMask16\",#\"AccFactor08\",\n    dataset_cache_file=\"dataset_cache_windows.pkl\"\n    )\n\nUsing dataset cache file\n\n\n\ntrain_dataset, test_dataset = random_split(dataset, (0.8, 0.2))\n\n\ntrain_dataloader = DataLoader(dataset=train_dataset, batch_size=1, shuffle=False)\ntest_dataloader  = DataLoader(dataset= test_dataset, batch_size=1, shuffle=False)\n\n\nx, y, mask = next(iter(train_dataloader))\n\n\nphysics = DynamicMRI((1, 2, 12, 512, 256))\ny2 = physics(x, mask=mask)\n\n\nx_hat = physics.A_adjoint(y2, mask=mask)\n\n\nmodel = ArtifactRemovalCRNN(\n    CRNN(num_cascades=2)\n).to(\"cpu\")\n\n\nphysics.set_mask(mask)\nx_recon = model(y, physics)\n\n\nplot_gif([x, y, y2, x_hat, mask], titles=[\"x\", \"y\", \"y2\", \"x_hat\", \"mask\"], display=True)\n\n\n\n\n\n\n\n\n  \n  \n    \n    \n      \n          \n      \n        \n      \n          \n      \n          \n      \n          \n      \n          \n      \n          \n      \n          \n      \n          \n    \n    \n      \n      Once\n      \n      Loop\n      \n      Reflect\n    \n  \n\n\n\n\n\n\n\n\n\n\nimport deepinv as dinv\n\n\nclass Undersampling(dinv.physics.DecomposablePhysics):\n    def __init__(self, mask, device=\"cpu\", **kwargs):\n        super().__init__(**kwargs)\n        self.mask = torch.nn.Parameter(mask.to(device), requires_grad=False)\n\n\nloss = YamanSplittingLoss(split_ratio=0.9)\n\nmask2 = loss.subsample_mask(loss.rng, physics.mask.data.detach().clone())\n\ninp = Undersampling(mask2, device=y.device)\n\ninp2 = Undersampling(mask - mask2, device=y.device)\n\nphysics1 = DynamicMRI(physics.img_size)\nphysics1.set_mask(mask2)\nphysics2 = DynamicMRI(physics.img_size)\nphysics2.set_mask(mask - mask2)\n\n# divide measurements\ny1 = inp.A(y)\ny2 = inp2.A(y)\n\n\nplot_videos([\n    mask, mask2, mask-mask2, \n    y, y1, y2, \n    physics.A_adjoint(y, mask=mask), physics1.A_adjoint(y1, mask=mask2), physics2.A_adjoint(y2, mask=mask-mask2)\n], display=True)\n\n\n\n\n\n\n\n\n  \n  \n    \n    \n      \n          \n      \n        \n      \n          \n      \n          \n      \n          \n      \n          \n      \n          \n      \n          \n      \n          \n    \n    \n      \n      Once\n      \n      Loop\n      \n      Reflect"
   },
+  {
+    "objectID": "demo_splitting_loss.html",
+    "href": "demo_splitting_loss.html",
+    "title": "Self-supervised learning with measurement splitting",
+    "section": "",
+    "text": "We demonstrate self-supervised learning with measurement splitting, to train a denoiser network on the MNIST dataset.\nMeasurement splitting constructs a ground-truth free loss \\(\\frac{m}{m_2}\\| y_2 - A_2 \\inversef{y_1}{A_1}\\|^2\\) by splitting the measurement and the forward operator using a randomly generated mask.\nSee :class:deepinv.loss.SplittingLoss for full details.\nimport deepinv as dinv\nfrom torch.utils.data import DataLoader\nimport torch\nfrom torchvision import transforms, datasets\nfrom deepinv.models.utils import get_weights_url\n\nc:\\Users\\s2558406\\Documents\\Repos\\deepinv\\venv\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n  from .autonotebook import tqdm as notebook_tqdm\ntorch.manual_seed(0)\ndevice = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else \"cpu\""
+  },
+  {
+    "objectID": "demo_splitting_loss.html#train-and-test-network",
+    "href": "demo_splitting_loss.html#train-and-test-network",
+    "title": "Self-supervised learning with measurement splitting",
+    "section": "Train and test network",
+    "text": "Train and test network\n\ntrainer = dinv.Trainer(\n    model=model,\n    physics=physics,\n    epochs=1,\n    losses=loss,\n    optimizer=optimizer,\n    device=device,\n    train_dataloader=train_dataloader,\n    plot_images=False,\n    save_path=None,\n    verbose=True,\n    show_progress_bar=False,\n    wandb_vis=False,\n)\n\nmodel = trainer.train()\n\nThe model has 444737 trainable parameters\nTrain epoch 0: TotalLoss=0.017, PSNR=11.59\n\n\nTest and visualise the model outputs using a small test set. We set the output to average over 50 iterations of random mask realisations. The trained model improves on the no-learning reconstruction by ~3dB.\n\ntrainer.plot_images = True\nmodel.MC_samples = 50\ntrainer.test(test_dataloader)\n\n\n\n\nTest PSNR: No learning rec.: 19.356+-1.523 | Model: 23.150+-1.996. \n\n\n(23.150262069702148,\n 1.9961603806415074,\n 19.355849266052246,\n 1.5233685706910973)\n\n\nSince this is a denoising example, above, we have set eval_split_output to True (see :class:deepinv.loss.SplittingLoss for details). Alternatively, we get worse results when we set eval_split_output to False:\n\nmodel.eval_split_output = False\ntrainer.test(test_dataloader)\n\n\n\n\nTest PSNR: No learning rec.: 19.356+-1.523 | Model: 14.441+-1.520. \n\n\n(14.44115686416626, 1.5199918435073472, 19.355849266052246, 1.5233685706910973)\n\n\nFurthermore, we can disable measurement splitting at evaluation altogether by setting eval_split_input to False (this is done in SSDU):\n\nmodel.eval_split_input = False\ntrainer.test(test_dataloader)\n\n\n\n\nTest PSNR: No learning rec.: 19.356+-1.523 | Model: 9.650+-1.670. \n\n\n(9.650477170944214, 1.6700495788637173, 19.355849266052246, 1.5233685706910973)"
+  },
   {
     "objectID": "demo_ei.html",
     "href": "demo_ei.html",
@@ -14,17 +28,17 @@
     "text": "Train a neural network to solve an image inpainting inverse problem, using perspective-EI and the deepinv library.\n\nimport deepinv as dinv\n\n\n\ntorch imports\nimport torch\nfrom torch.utils.data import DataLoader, random_split\nfrom torchvision.datasets import ImageFolder\nfrom torchvision.transforms import Compose, ToTensor, CenterCrop, Resize\nfrom torchvision.datasets.utils import download_and_extract_archive\n\n\nDefine inpainting experiment to reconstruct images from images masked with a random mask:\n\nphysics = dinv.physics.Inpainting((3, 256, 256), mask=0.6, device=\"cpu\")\n\nLoad Urban100 dataset of natural urban scenes:\n\n\nDownload dataset from HuggingFace\ndownload_and_extract_archive(\n    \"https://huggingface.co/datasets/eugenesiow/Urban100/resolve/main/data/Urban100_HR.tar.gz?download=true\",\n    \"Urban100\",\n    filename=\"Urban100_HR.tar.gz\",\n    md5=\"65d9d84a34b72c6f7ca1e26a12df1e4c\"\n)\n\n\n\ntrain_dataset, test_dataset = random_split(\n    ImageFolder(\"Urban100\", transform=Compose([ToTensor(), Resize(256), CenterCrop(256)])),\n    (0.8, 0.2)\n    )\n    \ntrain_dataloader, test_dataloader = DataLoader(train_dataset, shuffle=True), DataLoader(test_dataset)\n\nAs these scenes are imaged with a camera free to move and rotate in the world, we can impose perspective invariance on the unknown image set \\(x\\in X\\). Define measurement consistency and EI losses:\n\ntransform = dinv.transform.Homography(theta_max=10)\n\nlosses = [\n    dinv.loss.MCLoss(), \n    dinv.loss.EILoss(transform)\n]\n\nFor training, use a small UNet for the model with Adam optimizer:\n\nmodel = dinv.models.UNet(\n    in_channels=3, \n    out_channels=3,\n    scales=2,\n    circular_padding=True,\n    batch_norm=False\n)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-8)\n\nTrain the model using deepinv’s Trainer:\n\nmodel = dinv.Trainer(\n    model=model,\n    physics=physics,\n    online_measurements=True,\n    train_dataloader=train_dataloader,\n    eval_dataloader=test_dataloader,\n    epochs=1,\n    losses=losses,\n    optimizer=optimizer,\n    verbose=True,\n    show_progress_bar=False,\n    save_path=None,\n    device=\"cpu\"\n).train()\n\nThe model has 444867 trainable parameters\nEval epoch 0: PSNR=10.078\nTrain epoch 0: MCLoss=0.002, EILoss=0.021, TotalLoss=0.023, PSNR=15.948\n\n\nShow results of a pretrained model trained using a larger UNet for 40 epochs:\n\n\nLoad pretrained model from HuggingFace\nmodel = dinv.models.UNet(\n    in_channels=3, \n    out_channels=3,\n    scales=3,\n    circular_padding=True,\n    batch_norm=False\n)\n\nckpt = torch.hub.load_state_dict_from_url(\n    dinv.models.utils.get_weights_url(\"ei\", \"Urban100_inpainting_homography_model.pth\"),\n    map_location=\"cpu\",\n)\n\nmodel.load_state_dict(ckpt[\"state_dict\"])\n\n\n<All keys matched successfully>\n\n\n\nx, _ = next(iter(train_dataloader))\ny = physics(x)\nx_hat = model(y)\n\ndinv.utils.plot([x, y, x_hat], [\"x\", \"y\", \"reconstruction\"])"
   },
   {
-    "objectID": "demo_splitting_loss.html",
-    "href": "demo_splitting_loss.html",
+    "objectID": "demo_splitting_loss_tomography.html",
+    "href": "demo_splitting_loss_tomography.html",
     "title": "Self-supervised learning with measurement splitting",
     "section": "",
-    "text": "We demonstrate self-supervised learning with measurement splitting, to train a denoiser network on the MNIST dataset.\nMeasurement splitting constructs a ground-truth free loss \\(\\frac{m}{m_2}\\| y_2 - A_2 \\inversef{y_1}{A_1}\\|^2\\) by splitting the measurement and the forward operator using a randomly generated mask.\nSee :class:deepinv.loss.SplittingLoss for full details.\nimport deepinv as dinv\nfrom torch.utils.data import DataLoader\nimport torch\nfrom torchvision import transforms, datasets\nfrom deepinv.models.utils import get_weights_url\n\nc:\\Users\\s2558406\\Documents\\Repos\\deepinv\\venv\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n  from .autonotebook import tqdm as notebook_tqdm\ntorch.manual_seed(0)\ndevice = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else \"cpu\""
+    "text": "We demonstrate self-supervised learning with measurement splitting, to train a denoiser network on the MNIST dataset. The physics here is noisy computed tomography, as is the case in Noise2Inverse. Note this example can also be easily applied to undersampled multicoil MRI as is the case in SSDU.\nMeasurement splitting constructs a ground-truth free loss \\(\\frac{m}{m_2}\\| y_2 - A_2 \\inversef{y_1}{A_1}\\|^2\\) by splitting the measurement and the forward operator using a randomly generated mask.\nSee :class:deepinv.loss.SplittingLoss for full details.\nimport deepinv as dinv\nfrom torch.utils.data import DataLoader\nimport torch\nfrom torchvision import transforms, datasets\nfrom deepinv.models.utils import get_weights_url\n\nc:\\Users\\s2558406\\Documents\\Repos\\deepinv\\venv\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n  from .autonotebook import tqdm as notebook_tqdm\ntorch.manual_seed(0)\ndevice = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else \"cpu\""
   },
   {
-    "objectID": "demo_splitting_loss.html#train-and-test-network",
-    "href": "demo_splitting_loss.html#train-and-test-network",
+    "objectID": "demo_splitting_loss_tomography.html#train-and-test-network",
+    "href": "demo_splitting_loss_tomography.html#train-and-test-network",
     "title": "Self-supervised learning with measurement splitting",
     "section": "Train and test network",
-    "text": "Train and test network\n\ntrainer = dinv.Trainer(\n    model=model,\n    physics=physics,\n    epochs=1,\n    losses=loss,\n    optimizer=optimizer,\n    device=device,\n    train_dataloader=train_dataloader,\n    plot_images=False,\n    save_path=None,\n    verbose=True,\n    show_progress_bar=False,\n    wandb_vis=False,\n)\n\nmodel = trainer.train()\n\nThe model has 444737 trainable parameters\nTrain epoch 0: TotalLoss=0.017, PSNR=11.59\n\n\nTest and visualise the model outputs using a small test set. We set the output to average over 50 iterations of random mask realisations. The trained model improves on the no-learning reconstruction by ~3dB.\n\ntrainer.plot_images = True\nmodel.MC_samples = 50\ntrainer.test(test_dataloader)\n\n\n\n\nTest PSNR: No learning rec.: 19.356+-1.523 | Model: 23.150+-1.996. \n\n\n(23.150262069702148,\n 1.9961603806415074,\n 19.355849266052246,\n 1.5233685706910973)\n\n\nSince this is a denoising example, above, we have set eval_split_output to True (see :class:deepinv.loss.SplittingLoss for details). Alternatively, we get worse results when we set eval_split_output to False:\n\nmodel.eval_split_output = False\ntrainer.test(test_dataloader)\n\n\n\n\nTest PSNR: No learning rec.: 19.356+-1.523 | Model: 14.441+-1.520. \n\n\n(14.44115686416626, 1.5199918435073472, 19.355849266052246, 1.5233685706910973)\n\n\nFurthermore, we can disable measurement splitting at evaluation altogether by setting eval_split_input to False (this is done in SSDU):\n\nmodel.eval_split_input = False\ntrainer.test(test_dataloader)\n\n\n\n\nTest PSNR: No learning rec.: 19.356+-1.523 | Model: 9.650+-1.670. \n\n\n(9.650477170944214, 1.6700495788637173, 19.355849266052246, 1.5233685706910973)"
+    "text": "Train and test network\n\ntrainer = dinv.Trainer(\n    model=model,\n    physics=physics,\n    epochs=3,\n    losses=loss,\n    optimizer=optimizer,\n    device=device,\n    train_dataloader=train_dataloader,\n    plot_images=False,\n    save_path=None,\n    verbose=True,\n    show_progress_bar=False,\n    wandb_vis=False,\n)\n\nmodel = trainer.train()\n\nThe model has 444737 trainable parameters\nTrain epoch 0: TotalLoss=0.032, PSNR=29.155\nTrain epoch 1: TotalLoss=0.035, PSNR=28.895\nTrain epoch 2: TotalLoss=0.035, PSNR=28.837\n\n\nTest and visualise the model outputs using a small test set. We set the output to average over 5 iterations of random mask realisations. The trained model improves on the no-learning reconstruction by ~7dB.\n\ntrainer.plot_images = True\ntrainer.test(test_dataloader, pinv=True)\n\n\n\n\nTest PSNR: No learning rec.: 24.549+-1.052 | Model: 31.911+-2.220. \n\n\n(31.911066627502443, 2.219884675922773, 24.548791694641114, 1.0523162766060832)\n\n\nDemonstrate the effect of not averaging over multiple realisations of the splitting mask at evaluation time, by setting eval_n_samples=1. We have a worse performance:\n\nmodel.eval_n_samples = 1\ntrainer.test(test_dataloader, pinv=True)\n\n\n\n\nTest PSNR: No learning rec.: 24.549+-1.052 | Model: 30.434+-2.487. \n\n\n(30.434445762634276, 2.486670644991658, 24.548791694641114, 1.0523162766060832)\n\n\nFurthermore, we can disable measurement splitting at evaluation altogether by setting eval_split_input to False (this is done in SSDU). This generally is worse than MC averaging:\n\nmodel.eval_split_input = False\ntrainer.test(test_dataloader, pinv=True)\n\n\n\n\nTest PSNR: No learning rec.: 24.549+-1.052 | Model: 31.003+-2.107. \n\n\n(31.002875900268556, 2.106733038650352, 24.548791694641114, 1.0523162766060832)"
   }
 ]