Generalising Video Classification Training Script (references/video_classification/train.py)

[rsomani95]

The references/video_classification/train.py script can be generalised for finetuning on any dataset with just a few changes.

3 additional arguments would be needed:

parser.add_argument('--train-dir', default='train_avi-480p', help='name of train dir')
parser.add_argument('--val-dir', default='val_avi-480p', help='name of val dir')
parser.add_argument('--output-classes', default=None, help='no. of output classes (if finetuning)')

Some minimal changes to the script:

Modifying lines 119-120

traindir = os.path.join(args.data_path, args.train_dir)
valdir = os.path.join(args.data_path, args.valid_dir)

One additional line after line 205:

#line 205
model = torchvision.models.video.__dict__[args.model](pretrained=args.pretrained)

#additional chunk for fine-tuning
if args.output_classes is not None:
    model.fc.out_features = args.output_classes

#line 206
model.to(device)

If this makes sense to you guys, I'll be happy to put in a PR for it.

[rsomani95]

I sent in a PR (#1544) since it's easier to see the changes there

[fmassa]

Replied in the PR, thanks for the PR!

[datumbox]

This issue appears to be stale, so I'm closing it. If you think it should remain open, feel free to reopen it.