def test():
from common.utils.misc import print_dict
model = PlainCNN()
print(model)
data_batch = {
'image': torch.rand(4, 3, 128, 128),
'action': torch.randint(3, [4])
}
pd_dict = model(data_batch)
print_dict(pd_dict)
loss_dict = model.compute_losses(pd_dict, data_batch)
print_dict(loss_dict)
def test():
from common.utils.misc import print_dict
from torch_geometric.data import Data, Batch
model = EdgeConvNet()
print(model)
data = Data(
x=torch.rand(4, 3, 16, 16),
action=torch.randint(3, [1]),
pos=torch.rand(4, 4),
edge_index=torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]],
dtype=torch.int64),
size=torch.tensor([1], dtype=torch.int64),
)
data_batch = Batch.from_data_list([data])
pd_dict = model(data_batch)
print_dict(pd_dict)
loss_dict = model.compute_losses(pd_dict, data_batch)
print_dict(loss_dict)
def test():
from common.utils.misc import print_dict
model = SPACE_v1()
print(model)
data_batch = {
'image': torch.rand(4, 3, 64, 64),
'z_pres_p_prior': 0.1,
'z_where_loc_prior': torch.tensor([0.0, 0.0, 0.0, 0.0]),
'z_where_scale_prior': torch.tensor([0.2, 0.2, 0.2, 0.2]),
'z_what_loc_prior': 0.0,
'z_what_scale_prior': 1.0,
'z_depth_loc_prior': 0.0,
'z_depth_scale_prior': 1.0,
'fg_recon_scale_prior': 0.15,
'bg_recon_scale_prior': 0.15,
}
pd_dict = model(data_batch)
print_dict(pd_dict)
loss_dict = model.compute_losses(pd_dict, data_batch)
print_dict(loss_dict)