| 78 | |
| 79 | |
| 80 | def predict(cfg, task, dataset): |
| 81 | dataloader = data.DataLoader(dataset, cfg.batch_size, shuffle=False, num_workers=0) |
| 82 | device = torch.device(cfg.gpus[0]) |
| 83 | task = task.cuda(device) |
| 84 | task.eval() |
| 85 | seq_prob = [] |
| 86 | for batch in tqdm(dataloader): |
| 87 | batch = utils.cuda(batch, device=device) |
| 88 | with torch.no_grad(): |
| 89 | prob, sizes = task.inference(batch) |
| 90 | cum_sizes = sizes.cumsum(dim=0) |
| 91 | for i in range(len(sizes)): |
| 92 | seq_prob.append(prob[cum_sizes[i]-sizes[i]:cum_sizes[i]]) |
| 93 | return seq_prob |
| 94 | |
| 95 | |
| 96 | def get_prob(seq_prob, mutations, offsets): |