Change torchvision model by ArgumentParser


I would like to add an argument to my script in which the user can quickly decide for a torchvision model.

parser = argparse.ArgumentParser()
parser.add_argument('--arch', dest='torchmodel', action='store')
torchmodel = args.torchmodel

The model is set by, e.g. for vgg19:
model = models.vgg19(pretrained=True)

I would like make the models.attribute a placeholder, i.e. the string that is passed with “–arch ‘string’” is supposed to fill in the model-name in models.model_name(pretrained=True).

Is there a way to implement this? Thank you!

1 Like

Sure. Try

Instead of line 17, you load your model using an if...elif...else block.

Alternatively, use models.__dict__ to select the specified model as shown here.


Amazing, exactly what I was looking for! Thanks a lot!