Search code examples
pythonpython-3.xstringvariablesargs

How to change more simple?


import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--model1', default='resnet18', type=str)
parser.add_argument('--model2', default='resnet34', type=str)
parser.add_argument('--model3', default='resnet50', type=str)
parser.add_argument('--model4', default='resnet101', type=str)
parser.add_argument('--model_count', default=4, type=int)

args = parser.parse_args()

model1 = create_model(args.model1)
model2 = create_model(args.model2)
model3 = create_model(args.model3)
model4 = create_model(args.model4)

I want to change above code using loop, how can I do? I tried like below but it does not work.

for i in range(1, args.model_count+1):
    globals()["model{}".format(i)] = create_model(globals()["args.model{}".format(i)])

Traceback (most recent call last):
  globals()["model{}".format(i)] = create_model(globals()["args.model{}".format(i)])
KeyError: 'args.model1'

Solution

  • import argparse
    
    parser = argparse.ArgumentParser()
    args = [('--model1', 'resnet18'), ('--model2', 'resnet18'), 
        ('--model3', 'resnet50'), ('--model4', 'resnet101')]
    for model, net in args:
        parser.add_argument(model, default=net, type=str)
    parser.add_argument('--model_count', default=4, type=int)
    
    args = parser.parse_args()
    
    models_dict = {}
    for i in range(1, 5):
        models_dict[f'model{i}'] = create_model(getattr(args, f'model{i}')