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}" )