File size: 568 Bytes
beb7843 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | from .yowo.build import build_yowo
def build_model(args,
d_cfg,
m_cfg,
device,
num_classes=3,
trainable=False,
resume=None):
# build action detector
if 'yowo_v2_' in args.version:
model, criterion = build_yowo(
args=args,
d_cfg=d_cfg,
m_cfg=m_cfg,
device=device,
num_classes=num_classes,
trainable=trainable,
resume=resume
)
return model, criterion
|