def get_model(args):
# initiate model
input_channels = int(args.use_multiview) * 128 + int(
args.use_normal) * 3 + int(
args.use_color) * 3 + int(not args.no_height)
model = RefNet(num_class=DC.num_class,
num_heading_bin=DC.num_heading_bin,
num_size_cluster=DC.num_size_cluster,
mean_size_arr=DC.mean_size_arr,
input_feature_dim=input_channels,
num_proposal=args.num_proposals,
use_lang_classifier=(not args.no_lang_cls),
use_bidir=args.use_bidir,
no_reference=args.no_reference,
batch_size=args.batch_size)
# trainable model
if args.use_pretrained:
# load model
print("loading pretrained VoteNet...")
pretrained_model = RefNet(num_class=DC.num_class,
num_heading_bin=DC.num_heading_bin,
num_size_cluster=DC.num_size_cluster,
mean_size_arr=DC.mean_size_arr,
num_proposal=args.num_proposals,
input_feature_dim=input_channels,
use_bidir=args.use_bidir,
no_reference=True)
pretrained_path = os.path.join(CONF.PATH.OUTPUT, args.use_pretrained,
"model_last.pth")
pretrained_model.load_state_dict(torch.load(pretrained_path),
strict=False)
# mount
model.backbone_net = pretrained_model.backbone_net
model.vgen = pretrained_model.vgen
model.proposal = pretrained_model.proposal
if args.no_detection:
# freeze pointnet++ backbone
for param in model.backbone_net.parameters():
param.requires_grad = False
# freeze voting
for param in model.vgen.parameters():
param.requires_grad = False
# freeze detector
for param in model.proposal.parameters():
param.requires_grad = False
# to CUDA
model = model.cuda()
return model
def get_model(args):
# initiate model
input_channels = int(args.use_multiview) * 128 + int(args.use_normal) * 3 + int(args.use_color) * 3 + int(not args.no_height)
model = RefNet(
num_class=DC.num_class,
num_heading_bin=DC.num_heading_bin,
num_size_cluster=DC.num_size_cluster,
mean_size_arr=DC.mean_size_arr,
input_feature_dim=input_channels,
num_proposal=args.num_proposals,
use_lang_classifier=(not args.no_lang_cls),
use_bidir=args.use_bidir,
no_reference=args.no_reference,
attn=args.self_attn
)
# trainable model
if args.use_pretrained:
# load model
print("loading pretrained VoteNet...")
pretrained_model = RefNet(
num_class=DC.num_class,
num_heading_bin=DC.num_heading_bin,
num_size_cluster=DC.num_size_cluster,
mean_size_arr=DC.mean_size_arr,
num_proposal=args.num_proposals,
input_feature_dim=input_channels,
use_bidir=args.use_bidir,
no_reference=True,
attn=args.self_attn,
)
pretrained_path = os.path.join(args.use_pretrained, "model_last.pth")
pretrained_model.load_state_dict(torch.load(pretrained_path), strict=False)
# mount
model.backbone_net = pretrained_model.backbone_net
model.vgen = pretrained_model.vgen
model.proposal = pretrained_model.proposal
if args.no_detection:
# freeze pointnet++ backbone
for param in model.backbone_net.parameters():
param.requires_grad = False
# freeze voting
for param in model.vgen.parameters():
param.requires_grad = False
# freeze detector
for param in model.proposal.parameters():
param.requires_grad = False
# to CUDA
is_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if is_cuda else "cpu")
model = model.to(device)
devices = [int(x) for x in args.devices]
print("devices",devices, "torch.cuda.device_count()",torch.cuda.device_count())
model = nn.DataParallel(model, device_ids=devices)
# model = model.cuda()
return model