44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
# Script to evaluate pretrained models.
|
|
|
|
from argparse import ArgumentParser
|
|
import logging
|
|
import sys
|
|
|
|
import torch
|
|
|
|
from demucs import train, pretrained, evaluate
|
|
|
|
|
|
def main():
|
|
torch.set_num_threads(1)
|
|
logging.basicConfig(stream=sys.stderr, level=logging.INFO)
|
|
parser = ArgumentParser("tools.test_pretrained",
|
|
description="Evaluate pre-trained models or bags of models "
|
|
"on MusDB.")
|
|
pretrained.add_model_flags(parser)
|
|
parser.add_argument('overrides', nargs='*',
|
|
help='Extra overrides, e.g. test.shifts=2.')
|
|
args = parser.parse_args()
|
|
|
|
xp = train.main.get_xp(args.overrides)
|
|
with xp.enter():
|
|
solver = train.get_solver(xp.cfg)
|
|
|
|
model = pretrained.get_model_from_args(args)
|
|
solver.model = model.to(solver.device)
|
|
solver.model.eval()
|
|
|
|
with torch.no_grad():
|
|
results = evaluate.evaluate(solver, xp.cfg.test.sdr)
|
|
print(results)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|