def main(dataname):
model = CAE()
name = "CAE"
model_dict = torch.load('models/CAE_model', map_location='cpu')
model.load_state_dict(model_dict)
model.eval
EPOCH = 1
BATCH_SIZE_TRAIN = 49950
# dataname = "demonstrations/demo_00_02.pkl"
save_model_path = "models/" + name + "_model"
best_model_path = "models/" + name + "_best_model"
train_data = MotionData(dataname)
train_set = DataLoader(dataset=train_data,
batch_size=BATCH_SIZE_TRAIN,
shuffle=True)
for epoch in range(EPOCH):
epoch_loss = 0.0
for batch, x in enumerate(train_set):
# loss = model(x)
z = model.encoder(x).tolist()
print(len(z))
# print(epoch, loss.item())
data = np.asarray(z)
return data
class Model(object):
def __init__(self):
self.model = CAE()
model_dict = torch.load('models/CAE_best_model', map_location='cpu')
self.model.load_state_dict(model_dict)
self.model.eval
def decoder(self, z, s):
if abs(z[0][0]) < 0.01:
return [0.0] * 6
# z = np.asarray([z])
z = np.asarray(z)
# print(s.shape)
z_tensor = torch.FloatTensor(np.concatenate((z,s),axis=1))
# print(z_tensor.shape)
a_tensor = self.model.decoder(z_tensor)
return a_tensor.tolist()
def encoder(self, a, s):
x = np.concatenate((a,s),axis=1)
x_tensor = torch.FloatTensor(x)
z = self.model.encoder(x_tensor)
return z.tolist()