Add new model
January 21, 2023 ยท View on GitHub
DC-Bench includes most of the commonly used models to evaluate the performance of synthetic datasets.
But it's also easy to extend to add new models.
Step 1: Create the new model
You can create the new model in whatever way you like as long as it's a subclass of the PyTorch nn.Module
Step 2: Integrate
After you are done creating the model, put it under the networks folder. Then open network_utils.py in the networks folder and add your model similar to the following code
elif model_name == 'resnet18':
return ResNet18(channel=channel, num_classes=num_classes)
Step 3: Run the eval
Now you have integrated your new model, you can starting using the new model as
python evaluator/evaluator.py --method DC --model NEW_MODEL