Change torchvision model by ArgumentParser

vision
1

Hi,

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
2

Sure. Try https://github.com/psyec1/Lipreading-PyTorch/blob/master/main.py

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

3

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

2 Likes
4

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