def init_inference():
global model
global device
if args.model == 'resnet18':
model = models.resnet18()
model.fc = torch.nn.Linear(512, 3)
elif args.model == 'samplenet':
model = SampleNet()
elif args.model == 'simplenet':
model = SimpleNet()
else:
raise NotImplementedError()
model.eval()
#model.load_state_dict(torch.load(args.pretrained_model))
if args.trt_module:
from torch2trt import TRTModule
if args.trt_conversion:
model.load_state_dict(torch.load(args.pretrained_model))
model = model.cuda()
x = torch.ones((1, 3, 240, 320)).cuda()
from torch2trt import torch2trt
model_trt = torch2trt(model, [x],
max_batch_size=100,
fp16_mode=True)
#model_trt = torch2trt(model, [x], max_batch_size=100)
torch.save(model_trt.state_dict(), args.trt_model)
exit()
model_trt = TRTModule()
#model_trt.load_state_dict(torch.load('road_following_model_trt_half.pth'))
model_trt.load_state_dict(torch.load(args.trt_model))
model = model_trt.to(device)
else:
model.load_state_dict(torch.load(args.pretrained_model))
model = model.to(device)
def init_inference():
global device
if args.model == 'resnet18':
model = models.resnet18()
model.fc = torch.nn.Linear(512, 3)
elif args.model == 'samplenet':
model = SampleNet()
elif args.model == 'simplenet':
model = SimpleNet()
else:
raise NotImplementedError()
model.eval()
model.load_state_dict(torch.load(args.pretrained_model))
model = model.cuda()
x = torch.ones((1, 3, 240, 320)).cuda()
from torch2trt import torch2trt
#model_trt = torch2trt(model, [x], max_batch_size=100, fp16_mode=True)
model_trt = torch2trt(model, [x], max_batch_size=100)
torch.save(model_trt.state_dict(), args.trt_model)