126 lines
4.4 KiB
Python
126 lines
4.4 KiB
Python
# SPDX-License-Identifier: BSD-2-Clause
|
|
#
|
|
# Copyright (c) 2021, Pl@ntNet
|
|
# All rights reserved.
|
|
#
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
#
|
|
# 1. Redistributions of source code must retain the above copyright notice, this
|
|
# list of conditions and the following disclaimer.
|
|
#
|
|
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
# this list of conditions and the following disclaimer in the documentation
|
|
# and/or other materials provided with the distribution.
|
|
#
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
|
|
import torch
|
|
|
|
|
|
def add_all_parsers(parser):
|
|
_add_loss_parser(parser)
|
|
_add_training_parser(parser)
|
|
_add_model_parser(parser)
|
|
_add_hardware_parser(parser)
|
|
_add_misc_parser(parser)
|
|
|
|
|
|
def _add_loss_parser(parser):
|
|
group_loss = parser.add_argument_group("Loss parameters")
|
|
group_loss.add_argument(
|
|
"--mu", type=float, default=0.0001, help="weight decay parameter"
|
|
)
|
|
|
|
|
|
def _add_training_parser(parser):
|
|
group_training = parser.add_argument_group("Training parameters")
|
|
group_training.add_argument("--lr", type=float, help="learning rate to use")
|
|
group_training.add_argument(
|
|
"--batch_size", type=int, default=32, help="default is 32"
|
|
)
|
|
group_training.add_argument("--n_epochs", type=int)
|
|
group_training.add_argument("--pretrained", action="store_true")
|
|
group_training.add_argument("--image_size", type=int, default=256)
|
|
group_training.add_argument("--crop_size", type=int, default=224)
|
|
group_training.add_argument("--epoch_decay", nargs="+", type=int, default=[])
|
|
group_training.add_argument(
|
|
"--k",
|
|
nargs="+",
|
|
help="value of k for computing the topk loss and computing topk accuracy",
|
|
required=True,
|
|
type=int,
|
|
)
|
|
|
|
|
|
def _add_model_parser(parser):
|
|
group_model = parser.add_argument_group("Model parameters")
|
|
group_model.add_argument(
|
|
"--model",
|
|
choices=[
|
|
"resnet18",
|
|
"resnet34",
|
|
"resnet50",
|
|
"resnet101",
|
|
"resnet152",
|
|
"densenet121",
|
|
"densenet161",
|
|
"densenet169",
|
|
"densenet201",
|
|
"mobilenet_v2",
|
|
"inception_v3",
|
|
"alexnet",
|
|
"squeezenet",
|
|
"shufflenet",
|
|
"wide_resnet50_2",
|
|
"wide_resnet101_2",
|
|
"vgg11",
|
|
"mobilenet_v3_large",
|
|
"mobilenet_v3_small",
|
|
"inception_resnet_v2",
|
|
"inception_v4",
|
|
"efficientnet_b0",
|
|
"efficientnet_b1",
|
|
"efficientnet_b2",
|
|
"efficientnet_b3",
|
|
"efficientnet_b4",
|
|
"vit_base_patch16_224",
|
|
],
|
|
default="resnet50",
|
|
help="choose the model you want to train on",
|
|
)
|
|
|
|
|
|
def _add_hardware_parser(parser):
|
|
group_hardware = parser.add_argument_group("Hardware parameters")
|
|
group_hardware.add_argument(
|
|
"--use_gpu", type=int, choices=[0, 1], default=torch.cuda.is_available()
|
|
)
|
|
|
|
|
|
def _add_misc_parser(parser):
|
|
group_misc = parser.add_argument_group("Miscellaneous parameters")
|
|
group_misc.add_argument(
|
|
"--seed", type=int, help="set the seed for reproductible experiments"
|
|
)
|
|
group_misc.add_argument(
|
|
"--num_workers",
|
|
type=int,
|
|
default=4,
|
|
help="number of workers for the data loader. Default is one. You can bring it up. "
|
|
"If you have memory errors go back to one",
|
|
)
|
|
group_misc.add_argument(
|
|
"--root", help="location of the train val and test directories"
|
|
)
|
|
group_misc.add_argument("--save_name_xp", help="name of the saving file")
|