Python torch.set_grad_enabled() Examples

def preeval_batch(self, dataset):
        torch.set_grad_enabled(False)
        refs = {}
        cands = {}
        i = 0
        for batch_idx in range(len(dataset.corpus)):
            batch_s, batch_o_s, batch_f, batch_pf, batch_pb, _, targets, _, list_oovs, source_len, max_source_oov, w2fs = dataset.get_batch(batch_idx)
            decoded_outputs, lengths = self.model(batch_s, batch_o_s, max_source_oov, batch_f, batch_pf, batch_pb, source_len, w2fs=w2fs)
            for j in range(len(lengths)):
                i += 1
                ref = self.prepare_for_bleu(targets[j])
                refs[i] = [ref]
                out_seq = []
                for k in range(lengths[j]):
                    symbol = decoded_outputs[j][k].item()
                    if symbol < self.vocab.size:
                        out_seq.append(self.vocab.idx2word[symbol])
                    else:
                        out_seq.append(list_oovs[j][symbol-self.vocab.size])
                out = self.prepare_for_bleu(out_seq)
                cands[i] = out

                # if i % 2500 == 0:
                #     print("Percentages:  %.4f" % (i/float(dataset.len)))
        return cands, refs 

def _adapt(self, batch_samples, set_grad=True):
        """Performs one MAML inner step to update the policy.

        Args:
            batch_samples (MAMLTrajectoryBatch): Samples data for one
                task and one gradient step.
            set_grad (bool): if False, update policy parameters in-place.
                Else, allow taking gradient of functions of updated parameters
                with respect to pre-updated parameters.

        """
        # pylint: disable=protected-access
        loss = self._inner_algo._compute_loss(*batch_samples[1:])

        # Update policy parameters with one SGD step
        self._inner_optimizer.zero_grad()
        loss.backward(create_graph=set_grad)

        with torch.set_grad_enabled(set_grad):
            self._inner_optimizer.step() 

def forward(self, x):
        with torch.set_grad_enabled(self.finetune):
            x = self.model.conv1(x)
            x = self.model.bn1(x)
            x = self.model.relu(x)
            x = self.model.maxpool(x)

            x = self.model.layer1(x)
            x = self.model.layer2(x)
            x = self.model.layer3(x)
            x = self.model.layer4(x)

        if not self.spatial_context:
            x = F.avg_pool2d(x, kernel_size=7).view(x.size(0), x.size(1))
        if hasattr(self, 'context_transform'):
            x = self.context_transform(x)
        if hasattr(self, 'context_nonlinearity'):
            x = self.context_nonlinearity(x)
        return x 

def interpolate_penalty(self, network, input, penalty_amount=0.5):

        input = input.detach()
        input = input.requires_grad_()

        if len(input) != 2:
            raise ValueError('tuple of 2 inputs required to interpolate')
        inp1, inp2 = input

        try:
            epsilon = network.inputs.e.view(-1, 1, 1, 1)
        except AttributeError:
            raise ValueError('You must initiate a uniform random variable'
                             '`e` to use interpolation')
        mid_in = ((1. - epsilon) * inp1 + epsilon * inp2)
        mid_in.requires_grad_()

        with torch.set_grad_enabled(True):
            mid_out = network(mid_in)
        gradient = self._get_gradient(mid_in, mid_out)
        gradient = gradient.view(gradient.size()[0], -1)
        penalty = ((gradient.norm(2, dim=1) - 1.) ** 2).mean()

        return penalty_amount * penalty 

def contractive_penalty(self, network, input, penalty_amount=0.5):

        if penalty_amount == 0.:
            return

        if not isinstance(input, (list, tuple)):
            input = [input]

        input = [inp.detach() for inp in input]
        input = [inp.requires_grad_() for inp in input]

        with torch.set_grad_enabled(True):
            output = network(*input)
        gradient = self._get_gradient(input, output)
        gradient = gradient.view(gradient.size()[0], -1)
        penalty = (gradient ** 2).sum(1).mean()

        return penalty_amount * penalty 

def main_inference():
    print("Loading config...")
    opt = TestOptions().parse()
    print("Loading dataset...")
    dset = TVQADataset(opt, mode=opt.mode)
    print("Loading model...")
    model = STAGE(opt)
    model.to(opt.device)
    cudnn.benchmark = True
    strict_mode = not opt.no_strict
    model_path = os.path.join("results", opt.model_dir, "best_valid.pth")
    model.load_state_dict(torch.load(model_path), strict=strict_mode)
    model.eval()
    model.inference_mode = True
    torch.set_grad_enabled(False)
    print("Evaluation Starts:\n")
    predictions = inference(opt, dset, model)
    print("predictions {}".format(predictions.keys()))
    pred_path = model_path.replace("best_valid.pth",
                                   "{}_inference_predictions.json".format(opt.mode))
    save_json(predictions, pred_path) 

def sample(seed, num_bits, num_samples, samples_per_row, _log, output_path=None):
    torch.set_grad_enabled(False)

    if output_path is None:
        output_path = 'samples.png'

    torch.manual_seed(seed)
    np.random.seed(seed)

    device = set_device()

    _, _, (c, h, w) = get_train_valid_data()

    flow = create_flow(c, h, w).to(device)
    flow.eval()

    preprocess = Preprocess(num_bits)

    samples = flow.sample(num_samples)
    samples = preprocess.inverse(samples)

    save_image(samples.cpu(), output_path,
               nrow=samples_per_row,
               padding=0) 

def test():
    torch.set_grad_enabled(False)
    net.eval()
    loss_avg = 0.0
    correct = 0
    for batch_idx, (data, target) in enumerate(test_loader):
        data, target = data.cuda(), target.cuda()

        # forward
        output = net(data)
        loss = F.cross_entropy(output, target)

        # accuracy
        pred = output.data.max(1)[1]
        correct += pred.eq(target.data).sum().item()

        # test loss average
        loss_avg += loss.item()

    state['test_loss'] = loss_avg / len(test_loader)
    state['test_accuracy'] = correct / len(test_loader.dataset)
    torch.set_grad_enabled(True)


# Main loop 

def figure(self, batch_s, batch_o_s, batch_f, batch_pf, batch_pb, max_source_oov, source_len, list_oovs, w2fs, type,
               visual):
        torch.set_grad_enabled(False)
        decoded_outputs, lengths, self_matrix, soft = self.model(batch_s, batch_o_s, max_source_oov, batch_f, batch_pf, batch_pb,
                                              source_len, w2fs=w2fs, fig=True)
        length = lengths[0]
        output = []
        # print(decoded_outputs)
        for i in range(length):
            symbol = decoded_outputs[0][i].item()
            if symbol < self.vocab.size:
                output.append(self.vocab.idx2word[symbol])
            else:
                output.append(list_oovs[symbol-self.vocab.size])
        output = [i for i in output if i != '<PAD>' and i != '<EOS>' and i != '<SOS>']
        print(self_matrix.size(), soft.size())
        pos = [str(i) for i in batch_pf[0].cpu().tolist()]
        combine = []
        for j in range(len(pos)):
            combine.append(visual[j] + " : " + pos[j])
        self.showAttention(pos, combine, self_matrix.cpu(), 'self.png', 0)
        self.showAttention(type, output[19:25], soft[19:25].cpu(), 'type.png', 1)
        # return output 

def get_C_hat_transpose():
    torch.set_grad_enabled(False)
    probs = []
    net.eval()
    count = 0
    for batch_idx, (data, target) in enumerate(train_gold_deterministic_loader):
        # we subtract num_classes because we added num_classes to gold so we could identify which example is gold in train_phase2
        data, target = data.cuda(), (target - num_classes).cuda()
        count += target.shape[0]

        # forward
        output = net(data)
        pred = F.softmax(output, dim=1)
        probs.extend(list(pred.data.cpu().numpy()))

    probs = np.array(probs, dtype=np.float32)
    C_hat = np.zeros((num_classes, num_classes))
    for label in range(num_classes):
        indices = np.arange(len(train_data_gold.train_labels))[
            np.isclose(np.array(train_data_gold.train_labels) - num_classes, label)]
        C_hat[label] = np.mean(probs[indices], axis=0, keepdims=True)

    torch.set_grad_enabled(True)
    return C_hat.T.astype(np.float32) 

def test():
    torch.set_grad_enabled(False)
    net.eval()
    loss_avg = 0.0
    correct = 0
    for batch_idx, (data, target) in enumerate(test_loader):
        data, target = data.cuda(), target.cuda()

        # forward
        output = net(data)
        loss = F.cross_entropy(output, target)

        # accuracy
        pred = output.data.max(1)[1]
        correct += pred.eq(target.data).sum().item()

        # test loss average
        loss_avg += loss.item()

    state['test_loss'] = loss_avg / len(test_loader)
    state['test_accuracy'] = correct / len(test_loader.dataset)
    torch.set_grad_enabled(True)


# Main loop 

def do_one_epoch(self, epoch, episodes):
        mode = "train" if self.VAE.training else "val"
        epoch_loss, accuracy, steps = 0., 0., 0
        data_generator = self.generate_batch(episodes)
        for x_t in data_generator:
            with torch.set_grad_enabled(mode == 'train'):
                x_hat, mu, logvar = self.VAE(x_t)
                loss = self.loss_fn(x_t, x_hat, mu, logvar)

            if mode == "train":
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

            epoch_loss += loss.detach().item()
            steps += 1
        self.log_results(epoch, epoch_loss / steps, prefix=mode)
        if mode == "val":
            self.early_stopper(-epoch_loss / steps, self.encoder)

    #             xim = x_hat.detach().cpu().numpy()[0].transpose(1,2,0)
    #             self.wandb.log({"example_reconstruction": [self.wandb.Image(xim, caption="")]}) 

def do_one_epoch(self, epoch, episodes):
        mode = "train" if self.naff.training else "val"
        epoch_loss, accuracy, steps = 0., 0., 0
        data_generator = self.generate_batch(episodes)
        for x_t, x_tn in data_generator:
            with torch.set_grad_enabled(mode == 'train'):
                x_tn_hat = self.naff(x_t)
                loss = self.loss_fn(x_tn_hat, x_tn)

            if mode == "train":
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

            epoch_loss += loss.detach().item()
            steps += 1
        self.log_results(epoch, epoch_loss / steps, prefix=mode)
        if mode == "val":
            self.early_stopper(-epoch_loss / steps, self.encoder) 

def test_simple_inference(self):
            if not torch.cuda.is_available():
                import pytest

                pytest.skip('test requires GPU and torch+cuda')

            ori_grad_enabled = torch.is_grad_enabled()
            root_dir = os.path.dirname(os.path.dirname(__name__))
            model_config = os.path.join(
                root_dir, 'configs/mask_rcnn/mask_rcnn_r50_fpn_1x_coco.py')
            detector = MaskRCNNDetector(model_config)
            await detector.init()
            img_path = os.path.join(root_dir, 'demo/demo.jpg')
            bboxes, _ = await detector.apredict(img_path)
            self.assertTrue(bboxes)
            # asy inference detector will hack grad_enabled,
            # so restore here to avoid it to influence other tests
            torch.set_grad_enabled(ori_grad_enabled) 

def worker(in_queue: InQueue,
           out_queue: OutQueue,
           device: torch.device,
           grad_mode: bool,
           ) -> None:
    """The main loop of a worker thread."""
    torch.set_grad_enabled(grad_mode)

    with use_device(device):
        while True:
            task = in_queue.get()

            if task is None:
                break

            try:
                batch = task.compute()
            except Exception:
                exc_info = cast(ExcInfo, sys.exc_info())
                out_queue.put((False, exc_info))
                continue

            out_queue.put((True, (task, batch)))

    done = (False, None)
    out_queue.put(done) 

def worker(in_queue: InQueue,
           out_queue: OutQueue,
           device: torch.device,
           grad_mode: bool,
           ) -> None:
    """The main loop of a worker thread."""
    torch.set_grad_enabled(grad_mode)

    with use_device(device):
        while True:
            task = in_queue.get()

            if task is None:
                break

            try:
                batch = task.compute()
            except Exception:
                exc_info = cast(ExcInfo, sys.exc_info())
                out_queue.put((False, exc_info))
                continue

            out_queue.put((True, (task, batch)))

    done = (False, None)
    out_queue.put(done) 

def worker(in_queue: InQueue,
           out_queue: OutQueue,
           device: torch.device,
           grad_mode: bool,
           ) -> None:
    """The main loop of a worker thread."""
    torch.set_grad_enabled(grad_mode)

    with use_device(device):
        while True:
            task = in_queue.get()

            if task is None:
                break

            try:
                batch = task.compute()
            except Exception:
                exc_info = cast(ExcInfo, sys.exc_info())
                out_queue.put((False, exc_info))
                continue

            out_queue.put((True, (task, batch)))

    done = (False, None)
    out_queue.put(done) 

def run_model(net, loader, criterion, optimizer, train = True):
    running_loss = 0
    running_accuracy = 0

    # Set mode
    if train:
        net.train()
    else:
        net.eval()


    for i, (X, y) in enumerate(loader):
        # Pass to gpu or cpu
        X, y = X.to(device), y.to(device)

        # Zero the gradient
        optimizer.zero_grad()

        with torch.set_grad_enabled(train):
            output = net(X)
            _, pred = torch.max(output, 1)
            loss = criterion(output, y)

        # If on train backpropagate
        if train:
            loss.backward()
            optimizer.step()

        # Calculate stats
        running_loss += loss.item()
        running_accuracy += torch.sum(pred == y.detach())
    return running_loss / len(loader), running_accuracy.double() / len(loader.dataset) 

def test(self, path_to_image):
        # main loop
        torch.set_grad_enabled(False)
        print(f"test for image: {path_to_image}", flush=True)

        if self.is_cuda:
            model = self.model.cuda().eval()
        else:
            model = self.model.eval()

        img = cv2.imread(path_to_image)
        img = cv2.resize(img, (self.img_size, self.img_size))
        img_reverse = img[..., [2, 1, 0]]
        img = torch.from_numpy(img_reverse).float().permute(2, 0, 1).unsqueeze(0)

        if self.is_cuda:
            img = img.cuda()

        # measure elapsed time
        junc_pred, heatmap_pred, adj_mtx_pred = model(img)

        # visualize eval
        img = img.cpu().numpy()
        junctions_pred = junc_pred.cpu().numpy()
        adj_mtx = adj_mtx_pred.cpu().numpy()

        img_with_junc = draw_jucntions(img, junctions_pred)
        img_with_junc = img_with_junc[0].numpy()[None]
        img_with_junc = img_with_junc[:, ::-1, :, :]
        lines_pred, score_pred = graph2line(junctions_pred, adj_mtx)
        vis_line_pred = draw_lines(img_with_junc, lines_pred, score_pred)[0]
        vis_line_pred = vis_line_pred.permute(1, 2, 0).numpy()

        cv2.imshow("result", vis_line_pred) 

def train_epoches(t_dataset, v_dataset, model, n_epochs, teacher_forcing_ratio):
    eval_f = Evaluate_test()
    best_dev = 0
    train_loader = t_dataset.corpus
    len_batch = len(train_loader)
    epoch_examples_total = t_dataset.len
    for epoch in range(1, n_epochs + 1):
        model.train(True)
        torch.set_grad_enabled(True)
        epoch_loss = 0
        for batch_idx in range(len_batch):
            loss, num_examples = train_batch(t_dataset, batch_idx, model, teacher_forcing_ratio)
            epoch_loss += loss * num_examples
            sys.stdout.write(
                '%d batches processed. current batch loss: %f\r' %
                (batch_idx, loss)
            )
            sys.stdout.flush()
        epoch_loss /= epoch_examples_total
        log_msg = "Finished epoch %d with losses: %.4f" % (epoch, epoch_loss)
        print(log_msg)
        predictor = Predictor(model, v_dataset.vocab, args.cuda)
        print("Start Evaluating")
        cand, ref = predictor.preeval_batch(v_dataset)
        print('Result:')
        print('ref: ', ref[1][0])
        print('cand: ', cand[1])
        final_scores = eval_f.evaluate(live=True, cand=cand, ref=ref)
        epoch_score = 2*final_scores['ROUGE_L']*final_scores['Bleu_4']/(final_scores['Bleu_4']+ final_scores['ROUGE_L'])
        if epoch_score > best_dev:
            torch.save(model.state_dict(), args.save)
            print("model saved")
            best_dev = epoch_score 

def predict_file(self, dataset):
        torch.set_grad_enabled(False)
        i = 0
        lines = []
        for batch_idx in range(len(dataset.corpus)):
            batch_s, batch_o_s, batch_f, batch_pf, batch_pb, sources, targets, fields, list_oovs, source_len, \
                max_source_oov, w2fs = dataset.get_batch(batch_idx)
            decoded_outputs, lengths = self.model(batch_s, batch_o_s, max_source_oov, batch_f, batch_pf, batch_pb,
                                                  source_len, w2fs=w2fs)
            pos = batch_pf.tolist()
            for j in range(len(lengths)):
                line = {}
                line['sources'] = sources[j]
                line['fields'] = fields[j]
                i += 1
                line['refer'] = targets[j]
                line['pos'] = [p for p in pos[j] if p != 0]
                out_seq = []
                for k in range(lengths[j]):
                    symbol = decoded_outputs[j][k].item()
                    if symbol < self.vocab.size:
                        out_seq.append(self.vocab.idx2word[symbol])
                    else:
                        out_seq.append(list_oovs[j][symbol-self.vocab.size])
                line['output'] = out_seq
                self.overlap(line)

                if i % 2500 == 0:
                    print("Percentages:  %.4f" % (i/float(dataset.len)))
                lines.append(str(line)+'\n')
        return lines 

def do_one_epoch(self, epoch, episodes):
        mode = "train" if self.encoder.training and self.gru.training else "val"
        steps = 0
        step_losses = {i: [] for i in self.steps_gen()}
        step_accuracies = {i: [] for i in self.steps_gen()}

        data_generator = self.generate_batch(episodes)
        for sequence in data_generator:
            with torch.set_grad_enabled(mode == 'train'):
                sequence = sequence.to(self.device)
                sequence = sequence / 255.
                channels, w, h = self.config['obs_space'][-3:]
                flat_sequence = sequence.view(-1, channels, w, h)
                flat_latents = self.encoder(flat_sequence)
                latents = flat_latents.view(
                    self.batch_size, self.sequence_length, self.encoder.hidden_size)
                contexts, _ = self.gru(latents)
                loss = 0.
                for i in self.steps_gen():
                  predictions = self.discriminators[i](contexts[:, :-(i+1), :]).contiguous().view(-1, self.encoder.hidden_size)
                  targets = latents[:, i+1:, :].contiguous().view(-1, self.encoder.hidden_size)
                  logits = torch.matmul(predictions, targets.t())
                  step_loss = F.cross_entropy(logits, self.labels[i])
                  step_losses[i].append(step_loss.detach().item())
                  loss += step_loss

                  preds = torch.argmax(logits, dim=1)
                  step_accuracy = preds.eq(self.labels[i]).sum().float() / self.labels[i].numel()
                  step_accuracies[i].append(step_accuracy.detach().item())

            if mode == "train":
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

            steps += 1
        epoch_losses = {i: np.mean(step_losses[i]) for i in step_losses}
        epoch_accuracies = {i: np.mean(step_accuracies[i]) for i in step_accuracies}
        self.log_results(epoch, epoch_losses, epoch_accuracies, prefix=mode)
        if mode == "val":
            self.early_stopper(np.mean(list(epoch_accuracies.values())), self.encoder) 

def async_inference_detector(model, img):
    """Async inference image(s) with the detector.

    Args:
        model (nn.Module): The loaded detector.
        imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
            images.

    Returns:
        Awaitable detection results.
    """
    cfg = model.cfg
    device = next(model.parameters()).device  # model device
    # build the data pipeline
    test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
    test_pipeline = Compose(test_pipeline)
    # prepare data
    data = dict(img=img)
    data = test_pipeline(data)
    data = scatter(collate([data], samples_per_gpu=1), [device])[0]

    # We don't restore `torch.is_grad_enabled()` value during concurrent
    # inference since execution can overlap
    torch.set_grad_enabled(False)
    result = await model.aforward_test(rescale=True, **data)
    return result 

def predict(self, batch_s, batch_o_s, batch_f, batch_pf, batch_pb, max_source_oov, source_len, list_oovs, w2fs):
        torch.set_grad_enabled(False)
        decoded_outputs, lengths = self.model(batch_s, batch_o_s, max_source_oov, batch_f, batch_pf, batch_pb,
                                              source_len, w2fs=w2fs)
        length = lengths[0]
        output = []
        # print(decoded_outputs)
        for i in range(length):
            symbol = decoded_outputs[0][i].item()
            if symbol < self.vocab.size:
                output.append(self.vocab.idx2word[symbol])
            else:
                output.append(list_oovs[symbol-self.vocab.size])
        print(len(output))
        return ' '.join([i for i in output if i != '<PAD>' and i != '<EOS>' and i != '<SOS>']) 

def _compute_kl_constraint(self, all_samples, all_params, set_grad=True):
        """Compute KL divergence.

        For each task, compute the KL divergence between the old policy
        distribution and current policy distribution.

        Args:
            all_samples (list[list[MAMLTrajectoryBatch]]): Two
                dimensional list of MAMLTrajectoryBatch of size
                [meta_batch_size * (num_grad_updates + 1)]
            all_params (list[dict]): A list of named parameter dictionaries.
                Each dictionary contains key value pair of names (str) and
                parameters (torch.Tensor).
            set_grad (bool): Whether to enable gradient calculation or not.

        Returns:
            torch.Tensor: Calculated mean value of KL divergence.

        """
        theta = dict(self._policy.named_parameters())
        old_theta = dict(self._old_policy.named_parameters())

        kls = []
        for task_samples, task_params in zip(all_samples, all_params):
            for i in range(self._num_grad_updates):
                require_grad = i < self._num_grad_updates - 1 or set_grad
                self._adapt(task_samples[i], set_grad=require_grad)

            update_module_params(self._old_policy, task_params)
            with torch.set_grad_enabled(set_grad):
                # pylint: disable=protected-access
                kl = self._inner_algo._compute_kl_constraint(
                    task_samples[-1].observations)
            kls.append(kl)

            update_module_params(self._policy, theta)
            update_module_params(self._old_policy, old_theta)

        return torch.stack(kls).mean() 

def _compute_meta_loss(self, all_samples, all_params, set_grad=True):
        """Compute loss to meta-optimize.

        Args:
            all_samples (list[list[MAMLTrajectoryBatch]]): A two
                dimensional list of MAMLTrajectoryBatch of size
                [meta_batch_size * (num_grad_updates + 1)]
            all_params (list[dict]): A list of named parameter dictionaries.
                Each dictionary contains key value pair of names (str) and
                parameters (torch.Tensor).
            set_grad (bool): Whether to enable gradient calculation or not.

        Returns:
            torch.Tensor: Calculated mean value of loss.

        """
        theta = dict(self._policy.named_parameters())
        old_theta = dict(self._old_policy.named_parameters())

        losses = []
        for task_samples, task_params in zip(all_samples, all_params):
            for i in range(self._num_grad_updates):
                require_grad = i < self._num_grad_updates - 1 or set_grad
                self._adapt(task_samples[i], set_grad=require_grad)

            update_module_params(self._old_policy, task_params)
            with torch.set_grad_enabled(set_grad):
                # pylint: disable=protected-access
                last_update = task_samples[-1]
                loss = self._inner_algo._compute_loss(*last_update[1:])
            losses.append(loss)

            update_module_params(self._policy, theta)
            update_module_params(self._old_policy, old_theta)

        return torch.stack(losses).mean() 

def data_preprocess(self, *arg):
        output = []
        torch.set_grad_enabled(self._mode == 'Train')
        for data in arg:
            output.append(self._data_group_prepocess(data))

        output = tuple(output)
        if len(output) == 1:
            output = output[0]
        return output 

def makePrediction(config_obj, image_path, model, cuda=True, crf=False):
    torch.set_grad_enabled(False)

    image_id = extractIdFromPath(image_path)

    # Image preprocessing
    image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
    image = preProcessImage(image, config_obj, cuda)

    # Inference
    output = model(image)
    output = F.interpolate(
        output, size=image.shape[2:], mode="bilinear", align_corners=True
    )
    output = F.softmax(output, dim=1)
    output = output.data.cpu().numpy()[0]

    if crf:
        output = dense_crf(image_original, output)
    labelmap = np.argmax(output, axis=0)

    cocoResFormat = cocostuff.segmentationToCocoResult(labelmap, image_id)

    return cocoResFormat

# I cant create a JSON file as the segmentation is stores as bytes, need to decode 

def evaluate(adv=True):
    net.eval()
    if adv is False:
        torch.set_grad_enabled(False)
    running_loss = 0
    running_acc = 0
    count = 0
    for i, batch in enumerate(test_loader):
        bx = batch[0].cuda()
        by = batch[1].cuda()

        count += by.size(0)

        adv_bx = adversary(net, bx, by) if adv else bx
        with torch.no_grad():
            logits = net(adv_bx * 2 - 1)

        loss = F.cross_entropy(logits.data, by, reduction='sum')
        running_loss += loss.cpu().data.numpy()
        running_acc += (torch.max(logits, dim=1)[1] == by).float().sum(0).cpu().data.numpy()
    running_loss /= count
    running_acc /= count

    loss = running_loss
    acc = running_acc

    if adv is False:
        torch.set_grad_enabled(True)
    return loss, acc 

def worker(in_queue: InQueue,
           out_queue: OutQueue,
           device: torch.device,
           grad_mode: bool,
           ) -> None:
    """The main loop of a worker thread."""
    torch.set_grad_enabled(grad_mode)

    with use_device(device):
        while True:
            task = in_queue.get()

            if task is None:
                break

            try:
                batch = task.compute()
            except Exception:
                exc_info = cast(ExcInfo, sys.exc_info())
                out_queue.put((False, exc_info))
                continue

            out_queue.put((True, (task, batch)))

    done = (False, None)
    out_queue.put(done)