pouyashiri / lecapsnet

LECapsNet

This is the code for Light and Enhanced CapsNet (LE-CapsNet) published in ICMLA-2021.

The base code for this project is the following:
https://github.com/gram-ai/capsule-networks

Installation

Step 1. Install PyTorch and Torchvision:
conda install -c pytorch pytorch torchvision
Step 2. Install Torchnet:
pip install torchnet

Usage

The "Main.py" file trains the network and prints the results to the files in the specified folder (input args).
Parameters:
--dset: Choice of dataset (options: MNIST, F-MNIST, SVHN and CIFAR-10)
--nc: Number of classes in the chosen dataset
--w : The width/height of input images
--bsize: Batch size
--ne: Number of epochs to train the model
--niter: Number of iterations for DR algorithm
--fck: Fully-Connected Kernel size (K parameter of the CFC layer)
--fdim: The output dimensionality (D parameter of the CFC layer)
--ich: number of channels in the input image
--dec_type: The type of decoder used (options: FC, DECONV)
--res_folder: The output folder to print the results into
--albm: Whether or not use albumentation augmentations
--nc_recon: Performing the reconstruction in a single channel or all channels (options: 1,3)
--hard: Perform hard-training at the end or not (hard-training: training while tightening the bounds of the margin loss, options: 0,1)
--test_only: The option to only test the network against the test set (needs a checkpoint)
--checkpoint: The file address of the checkpoint file (used for hard training)