def __init__(self, env):
super().__init__(env)
# Assumes raycast observation space of len 20
self.observation_space = gym.spaces.Discrete(5)
self.v = VAE.load_from_checkpoint(
"/home/denpak/logs/vae_model/last.ckpt")
for param in self.v.parameters():
param.requires_grad = False
def _data_loader(self,
dataset: Dataset,
shuffle: bool = False) -> DataLoader:
return DataLoader(dataset,
batch_size=self.batch_size,
shuffle=shuffle,
num_workers=self.num_workers,
drop_last=self.drop_last,
pin_memory=self.pin_memory)
#model = VAE(64,latent_dim=5)
#trainer = pl.Trainer(gpus=1)
data = FishDataModule()
#trainer.fit(model,data)
v = VAE.load_from_checkpoint(
"/home/denpak/Research/BBEEncoders/vae_model/last.ckpt")
d = FishDataset(data_dir="/home/denpak/Research/BBEProject/Recordings",
labels_file="labels.csv")
x = torch.unsqueeze(d[0][0], 0)
x = v.encoder(x)
print(v.fc_mu(x)[0])
print(v.fc_var(x)[0])
x = torch.unsqueeze(d[1][0], 0)
x = v.encoder(x)
print(v.fc_mu(x)[0].detach().numpy())
print(v.fc_var(x)[0])
