def create_model(use_cuda):
tcn = define_model(use_cuda)
# tcn = PosNet()
if arguments.load_model:
model_path = os.path.join(arguments.model_folder, arguments.load_model)
# map_location allows us to load models trained on cuda to cpu.
tcn.load_state_dict(
torch.load(model_path, map_location=lambda storage, loc: storage))
if use_cuda:
tcn = tcn.cuda()
return tcn
def load_model(model_path, use_cuda=True):
tcn = define_model(use_cuda)
tcn = torch.nn.DataParallel(tcn, device_ids=range(torch.cuda.device_count()))
# tcn = PosNet()
# Change dict names if model was created with nn.DataParallel
state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
new_state_dict = OrderedDict()
# for k, v in state_dict.items():
# name = k[7:] # remove module.
# new_state_dict[name] = v
# map_location allows us to load models trained on cuda to cpu.
tcn.load_state_dict(state_dict)
if use_cuda:
tcn = tcn.cuda()
return tcn
self._kuka.applyAction(action)
bullet.stepSimulation()
if self._renders:
time.sleep(self._timeStep)
reward = self._reward()
frames = self.reward_frames()
self._envStepCounter += 1
done = self._termination()
if not (done):
self._observation = self.getExtendedObservation()
return self._observation, reward, done, frames
if __name__ == "__main__":
frame_size = (299, 299)
tcn = define_model(False)
model_path = os.path.join("./trained_models/tcn",
"inception-epoch-2000.pk")
tcn.load_state_dict(
torch.load(model_path, map_location=lambda storage, loc: storage))
env = ImitationEnv(renders=True,
video_dir='./data/validation/angle1',
tcn=tcn,
frame_size=frame_size)
done = False
while not (done):
observations, reward, done, _ = env.step(env.action_space.sample())
assert len(x.size()) == 4
def load_model(self, model_path, model):
tcn = define_model(self.use_cuda)
model_path = os.path.join(model_path, model)
tcn.load_state_dict(
torch.load(model_path, map_location=lambda storage, loc: storage))
return tcn
def load_model(model_path):
tcn = define_model(pretrained=False)
tcn.load_state_dict(
torch.load(model_path, map_location=lambda storage, loc: storage))
return tcn