def compute(
self,
model: Model,
graph: Graph,
trajectories: Trajectories,
pairwise_features: torch.Tensor,
):
"""Update the metrics for all trajectories in `trajectories`"""
self.init_metrics()
config = self.config
with torch.no_grad():
for trajectory_idx in tqdm(range(len(trajectories))):
observations = trajectories[trajectory_idx]
number_steps = None
if config.rw_edge_weight_see_number_step or config.rw_expected_steps:
if config.use_shortest_path_distance:
number_steps = (trajectories.leg_shortest_lengths(
trajectory_idx).float() * 1.1).long()
else:
number_steps = trajectories.leg_lengths(trajectory_idx)
observed, starts, targets = generate_masks(
trajectory_length=observations.shape[0],
number_observations=config.number_observations,
predict=config.target_prediction,
with_interpolation=config.with_interpolation,
device=config.device,
)
diffusion_graph = (graph if not config.diffusion_self_loops
else graph.add_self_loops())
predictions, _, rw_weights = model(
observations,
graph,
diffusion_graph,
observed=observed,
starts=starts,
targets=targets,
pairwise_node_features=pairwise_features,
number_steps=number_steps,
)
self.update_metrics(
trajectories,
graph,
observations,
observed,
starts,
targets,
predictions,
rw_weights,
trajectory_idx,
model.rw_non_backtracking,
)
def forward(self, source_seq, target_seq, source_lens, target_lens, pad_idx, train=False):
with torch.no_grad():
hidden_states, _ = self.lm.lm(self.lm.embed(source_seq))
logits = self.probe(hidden_states)
source_masks = -1000 * utils.generate_masks(source_lens).unsqueeze(2)
if train:
loss = F.binary_cross_entropy_with_logits((logits + source_masks).flatten(), target_seq.flatten())
return loss
else:
predictions = (torch.sigmoid(logits) >= 0.5).int()
return predictions + source_masks
import torch.nn as nn
import torch
import scipy.io as scio
import datetime
import os
import numpy as np
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
if not torch.cuda.is_available():
raise Exception('NO GPU!')
data_path = "./train"
test_path1 = "./test" # simulation data for comparison
mask, mask_s = generate_masks(data_path)
last_train = 100
model_save_filename = 'save_model'
block_size = 256
compress_rate = 8
first_frame_net = cnn1(compress_rate + 1).cuda()
rnn1 = forward_rnn().cuda()
rnn2 = backrnn().cuda()
# load pretrained model
if last_train != 0:
first_frame_net = torch.load(
'./model/' + model_save_filename +
"/first_frame_net_model_epoch_{}.pth".format(last_train))
pin_memory=True)
limit_set.setmode("image")
model.setmode("image")
counts = inference_image(limit_loader, model, device, mode='test')[1]
img_indices = np.select([counts != 0], [counts]).nonzero()[0]
indices = [i for i, g in enumerate(groups) if g in img_indices]
tiles = tiles[indices]
groups = groups[indices]
pseudo_masks = generate_masks(dataset,
tiles,
groups,
preprocess=args.preprocess,
output_path=os.path.join(
training_data_path,
args.pseudomask_dir))
trainset = Maskset(os.path.join(training_data_path, "training.h5"),
pseudo_masks,
augment=args.augment,
num_of_imgs=100 if args.debug else 0)
generate_masks(dataset,
tiles,
groups,
preprocess=args.preprocess,
output_path=os.path.join(training_data_path,
args.pseudomask_dir))
trainset = Maskset(os.path.join(training_data_path, "training.h5"),
def train_epoch(
model: Model,
graph: Graph,
optimizer: torch.optim.Optimizer,
config: Config,
train_trajectories: Trajectories,
pairwise_node_features: torch.Tensor,
):
"""One epoch of training"""
model.train()
print_cum_loss = 0.0
print_num_preds = 0
print_time = time.time()
print_every = len(
train_trajectories) // config.batch_size // config.print_per_epoch
trajectories_shuffle_indices = np.arange(len(train_trajectories))
if config.shuffle_samples:
np.random.shuffle(trajectories_shuffle_indices)
for iteration, batch_start in enumerate(
range(0,
len(trajectories_shuffle_indices) - config.batch_size + 1,
config.batch_size)):
optimizer.zero_grad()
loss = torch.tensor(0.0, device=config.device)
for i in range(batch_start, batch_start + config.batch_size):
trajectory_idx = trajectories_shuffle_indices[i]
observations = train_trajectories[trajectory_idx]
length = train_trajectories.lengths[trajectory_idx]
number_steps = None
if config.rw_edge_weight_see_number_step or config.rw_expected_steps:
if config.use_shortest_path_distance:
number_steps = (train_trajectories.leg_shortest_lengths(
trajectory_idx).float() * 1.1).long()
else:
number_steps = train_trajectories.leg_lengths(
trajectory_idx)
observed, starts, targets = generate_masks(
trajectory_length=observations.shape[0],
number_observations=config.number_observations,
predict=config.target_prediction,
with_interpolation=config.with_interpolation,
device=config.device,
)
diffusion_graph = graph if not config.diffusion_self_loops else graph.add_self_loops(
)
predictions, potentials, rw_weights = model(
observations,
graph,
diffusion_graph,
observed=observed,
starts=starts,
targets=targets,
pairwise_node_features=pairwise_node_features,
number_steps=number_steps,
)
print_num_preds += starts.shape[0]
l = (compute_loss(
config.loss,
train_trajectories,
observations,
predictions,
starts,
targets,
rw_weights,
trajectory_idx,
) / starts.shape[0])
loss += l
loss /= config.batch_size
print_cum_loss += loss.item()
loss.backward()
optimizer.step()
if (iteration + 1) % print_every == 0:
print_loss = print_cum_loss / print_every
print_loss /= print_num_preds
pred_per_second = 1.0 * print_num_preds / \
(time.time() - print_time)
print_cum_loss = 0.0
print_num_preds = 0
print_time = time.time()
progress_percent = int(100.0 * ((iteration + 1) // print_every) /
config.print_per_epoch)
print(
f"Progress {progress_percent}% | iter {iteration} | {pred_per_second:.1f} pred/s | loss {config.loss} {print_loss}"
)
