args = pickle.load(open(meta_file, 'rb'))['args']
args.cuda = not args_eval.no_cuda and torch.cuda.is_available()
args.batch_size = 100
args.dataset = args_eval.dataset
args.seed = 0
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
device = torch.device('cuda' if args.cuda else 'cpu')
dataset = utils.PathDataset(hdf5_file=args.dataset,
path_length=args_eval.num_steps)
eval_loader = data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
# Get data sample
obs = eval_loader.__iter__().next()[0]
input_shape = obs[0][0].size()
model = modules.ContrastiveSWM(embedding_dim=args.embedding_dim,
hidden_dim=args.hidden_dim,
action_dim=args.action_dim,
input_dims=input_shape,
num_objects=args.num_objects,
sigma=args.sigma,
decoder_model_file = os.path.join(args_eval.save_folder, 'decoder.pt')
args = pickle.load(open(meta_file, 'rb'))['args']
args.dataset = args_eval.dataset
args.cuda = not args_eval.no_cuda and torch.cuda.is_available()
args.batch_size = 100
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
device = torch.device('cuda' if args.cuda else 'cpu')
dataset = utils.PathDataset(hdf5_file=args.dataset, path_length=10)
def worker_init_fc(worker_index):
# make sure each worker has a different starting random seed
seed = np.random.get_state()[1][0]
np.random.seed(seed + worker_index)
train_loader = data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
worker_init_fn=worker_init_fc)
# Get data sample
args = pickle.load(open(meta_file, 'rb'))['args']
args.cuda = not args_eval.no_cuda and torch.cuda.is_available()
args.batch_size = 100
args.dataset = args_eval.dataset
args.seed = 0
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
device = torch.device('cuda' if args.cuda else 'cpu')
dataset = utils.PathDataset(hdf5_file=args.dataset,
action_encoding=args.action_encoding,
path_length=args_eval.num_steps,
truncate=50)
eval_loader = data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
# Get data sample
obs = eval_loader.__iter__().next()[0]
input_shape = obs[0][0].size()
model = modules.ContrastiveSWM(embedding_dim=args.embedding_dim,
hidden_dim=args.hidden_dim,
action_dim=args.action_dim,
input_dims=input_shape,
num_objects=args.num_objects,
def eval_c_swm(args, model):
torch.backends.cudnn.deterministic = True
# meta_file = os.path.join(args.save_folder, 'metadata.pkl')
# model_file = os.path.join(args.save_folder, 'model.pt')
# args = pickle.load(open(meta_file, 'rb'))['args']
args.cuda = not args.no_cuda and torch.cuda.is_available()
batch_size = 100
args.seed = 0
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
device = torch.device('cuda' if args.cuda else 'cpu')
dataset = utils.PathDataset(hdf5_file=args.dataset_eval,
path_length=args.num_steps)
eval_loader = data.DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=args.num_workers)
# Get data sample
obs = eval_loader.__iter__().next()[0]
input_shape = obs[0][0].size()
# model = modules.ContrastiveSWM(
# embedding_dim=args.embedding_dim,
# hidden_dim=args.hidden_dim,
# action_dim=args.action_dim,
# input_dims=input_shape,
# num_objects=args.num_objects,
# sigma=args.sigma,
# hinge=args.hinge,
# ignore_action=args.ignore_action,
# copy_action=args.copy_action,
# encoder=args.encoder).to(device)
# model.load_state_dict(torch.load(model_file))
model.eval()
# topk = [1, 5, 10]
topk = [1]
hits_at = defaultdict(int)
num_samples = 0
rr_sum = 0
pred_states = []
next_states = []
with torch.no_grad():
for batch_idx, data_batch in enumerate(eval_loader):
data_batch = [[t.to(device) for t in tensor]
for tensor in data_batch]
observations, actions = data_batch
if observations[0].size(0) != batch_size:
continue
obs = observations[0]
next_obs = observations[-1]
state = model.obj_encoder(model.obj_extractor(obs))
next_state = model.obj_encoder(model.obj_extractor(next_obs))
pred_state = state
for i in range(args.num_steps):
pred_trans = model.transition_model(pred_state, actions[i])
pred_state = pred_state + pred_trans
pred_states.append(pred_state.cpu())
next_states.append(next_state.cpu())
pred_state_cat = torch.cat(pred_states, dim=0)
next_state_cat = torch.cat(next_states, dim=0)
full_size = pred_state_cat.size(0)
# Flatten object/feature dimensions
next_state_flat = next_state_cat.view(full_size, -1)
pred_state_flat = pred_state_cat.view(full_size, -1)
dist_matrix = utils.pairwise_distance_matrix(next_state_flat,
pred_state_flat)
dist_matrix_diag = torch.diag(dist_matrix).unsqueeze(-1)
dist_matrix_augmented = torch.cat([dist_matrix_diag, dist_matrix],
dim=1)
# Workaround to get a stable sort in numpy.
dist_np = dist_matrix_augmented.numpy()
indices = []
for row in dist_np:
keys = (np.arange(len(row)), row)
indices.append(np.lexsort(keys))
indices = np.stack(indices, axis=0)
indices = torch.from_numpy(indices).long()
print('Processed {} batches of size {}'.format(batch_idx + 1,
batch_size))
labels = torch.zeros(indices.size(0),
device=indices.device,
dtype=torch.int64).unsqueeze(-1)
num_samples += full_size
print('Size of current topk evaluation batch: {}'.format(full_size))
for k in topk:
match = indices[:, :k] == labels
num_matches = match.sum()
hits_at[k] += num_matches.item()
match = indices == labels
_, ranks = match.max(1)
reciprocal_ranks = torch.reciprocal(ranks.double() + 1)
rr_sum += reciprocal_ranks.sum()
pred_states = []
next_states = []
for k in topk:
print('Hits @ {}: {}'.format(k, hits_at[k] / float(num_samples)))
print('MRR: {}'.format(rr_sum / float(num_samples)))
results_dict = {}
results_dict['H1'] = hits_at[1] / float(num_samples)
results_dict['MRR'] = rr_sum / float(num_samples)
return results_dict
