def get_model(args):
assert args.type in ['2d', '3d', 's3d']
if args.type == '2d':
print('Loading 2D-ResNet-152 ...')
model = models.resnet152(pretrained=True)
model = nn.Sequential(*list(model.children())[:-2], GlobalAvgPool())
model = model.cuda()
elif args.type == '3d':
print('Loading 3D-ResneXt-101 ...')
model = resnext.resnet101(
num_classes=400,
shortcut_type='B',
cardinality=32,
sample_size=112,
sample_duration=16,
last_fc=False)
model = model.cuda()
model_data = th.load(args.resnext101_model_path)
model.load_state_dict(model_data)
else:
print('Loading S3D ...')
model = S3D(
'model/s3d_dict.npy',
num_classes=512
)
model = model.cuda()
model_data = th.load(args.s3d_model_path)
model.load_state_dict(model_data)
# device = th.device('cuda:0')
# model.to(device)
model.eval()
print('loaded')
return model
def get_model(args):
assert args.type in ["2d", "3d", "ig"]
if args.type == "2d":
print("Loading 2D-ResNet-152 ...")
model = models.resnet152(pretrained=True)
model = nn.Sequential(*list(model.children())[:-2], GlobalAvgPool())
model = model.cuda()
elif args.type == "3d":
print("Loading 3D-ResneXt-101 ...")
model = resnext.resnet101(
num_classes=400,
shortcut_type="B",
cardinality=32,
sample_size=112,
sample_duration=16,
last_fc=False,
)
model = model.cuda()
model_data = th.load(args.resnext101_model_path)
model.load_state_dict(model_data)
else:
model = r2plus1d_34()
checkpoint = th.load(args.ig_model_path, map_location="cpu")
model.load_state_dict(checkpoint["model"])
model = th.nn.Sequential(*list(model.children())[:-2], AdaptivePool())
model = model.cuda()
model.eval()
print("loaded")
return model
def get_model(args):
assert args.type in ['2d', '3d', 'vmz', 's3d', 'vae']
if args.type == '2d':
print('Loading 2D-ResNet-152 ...')
import torchvision.models as models
model = models.resnet152(pretrained=True)
model = nn.Sequential(*list(model.children())[:-2], GlobalAvgPool())
model = model.cuda()
elif args.type == 'vmz':
print('Loading VMZ ...')
from vmz34 import r2plus1d_34
model = r2plus1d_34(pretrained_path=args.vmz_model_path,
pretrained_num_classes=487)
model = model.cuda()
elif args.type == 's3d':
# we use one copy of s3d instead of dup another one for feature extraction.
from mmpt.processors.models.s3dg import S3D
model = S3D('pretrained_models/s3d_dict.npy', 512)
model.load_state_dict(th.load('pretrained_models/s3d_howto100m.pth'))
model = model.cuda()
elif args.type == '3d':
print('Loading 3D-ResneXt-101 ...')
from videocnn.models import resnext
model = resnext.resnet101(num_classes=400,
shortcut_type='B',
cardinality=32,
sample_size=112,
sample_duration=16,
last_fc=False)
model = model.cuda()
model_data = th.load(args.resnext101_model_path)
model.load_state_dict(model_data)
elif args.type == 'vae':
from openaivae import OpenAIParallelDiscreteVAE
model = OpenAIParallelDiscreteVAE()
model = model.cuda()
else:
raise ValueError("model not supported yet.")
model.eval()
print('loaded')
return model
def get_model(args):
assert args.type in ['2d', '3d']
if args.type == '2d':
print('Loading 2D-ResNet-152 ...')
model = models.resnet152(pretrained=True)
model = nn.Sequential(*list(model.children())[:-2], GlobalAvgPool())
if th.cuda.is_available():
model = model.cuda()
else:
print('Loading 3D-ResneXt-101 ...')
model = resnext.resnet101(num_classes=400,
shortcut_type='B',
cardinality=32,
sample_size=112,
sample_duration=16,
last_fc=False)
if th.cuda.is_available():
model = model.cuda()
model_data = th.load(args.resnext101_model_path)
model.load_state_dict(model_data)
model.eval()
print('loaded')
return model