Hello,

When I load a stored checkpoint, I get the following error:

RuntimeError: output with shape [128, 3, 1] doesn't match the broadcast shape [128, 3, 3]

If I am reading a state_dict correctly, then I think there is probably a bug in you load_state_dict. For your convenience I slightly modified your CompressAI/examples/train.py example to accept also a checkpoint as input to continue from a previously stored checkpoint. For that, you just need to run the following code twice (I used bmshj2018-hyperprior model):

1. Once without --checkpoint-file for 1-2 epochs just to save a checkpoint
2. Then run the script with --checkpoint-file [/address/of/stored/checkpoint].

# Copyright 2020 InterDigital Communications, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import math
import random
import shutil
import sys

import torch
import torch.nn as nn
import torch.optim as optim

from torch.utils.data import DataLoader
from torchvision import transforms

from compressai.datasets import ImageFolder
from compressai.zoo import models


class RateDistortionLoss(nn.Module):
    """Custom rate distortion loss with a Lagrangian parameter."""

    def __init__(self, lmbda=1e-2):
        super().__init__()
        self.mse = nn.MSELoss()
        self.lmbda = lmbda

    def forward(self, output, target):
        N, _, H, W = target.size()
        out = {}
        num_pixels = N * H * W

        out["bpp_loss"] = sum(
            (torch.log(likelihoods).sum() / (-math.log(2) * num_pixels))
            for likelihoods in output["likelihoods"].values()
        )
        out["mse_loss"] = self.mse(output["x_hat"], target)
        out["loss"] = self.lmbda * 255 ** 2 * out["mse_loss"] + out["bpp_loss"]

        return out


class AverageMeter:
    """Compute running average."""

    def __init__(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


class CustomDataParallel(nn.DataParallel):
    """Custom DataParallel to access the module methods."""

    def __getattr__(self, key):
        try:
            return super().__getattr__(key)
        except AttributeError:
            return getattr(self.module, key)


def configure_optimizers(net, args):
    """Separate parameters for the main optimizer and the auxiliary optimizer.
    Return two optimizers"""

    parameters = set(
        p for n, p in net.named_parameters() if not n.endswith(".quantiles")
    )
    aux_parameters = set(
        p for n, p in net.named_parameters() if n.endswith(".quantiles")
    )

    # Make sure we don't have an intersection of parameters
    params_dict = dict(net.named_parameters())
    inter_params = parameters & aux_parameters
    union_params = parameters | aux_parameters

    assert len(inter_params) == 0
    assert len(union_params) - len(params_dict.keys()) == 0

    optimizer = optim.Adam(
        (p for p in parameters if p.requires_grad),
        lr=args.learning_rate,
    )
    aux_optimizer = optim.Adam(
        (p for p in aux_parameters if p.requires_grad),
        lr=args.aux_learning_rate,
    )
    return optimizer, aux_optimizer


def train_one_epoch(
    model, criterion, train_dataloader, optimizer, aux_optimizer, epoch, clip_max_norm
):
    model.train()
    device = next(model.parameters()).device

    for i, d in enumerate(train_dataloader):
        d = d.to(device)

        optimizer.zero_grad()
        aux_optimizer.zero_grad()

        out_net = model(d)

        out_criterion = criterion(out_net, d)
        out_criterion["loss"].backward()
        if clip_max_norm > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_max_norm)
        optimizer.step()

        aux_loss = model.aux_loss()
        aux_loss.backward()
        aux_optimizer.step()

        if i % 10 == 0:
            print(
                f"Train epoch {epoch}: ["
                f"{i*len(d)}/{len(train_dataloader.dataset)}"
                f" ({100. * i / len(train_dataloader):.0f}%)]"
                f'\tLoss: {out_criterion["loss"].item():.3f} |'
                f'\tMSE loss: {out_criterion["mse_loss"].item():.3f} |'
                f'\tBpp loss: {out_criterion["bpp_loss"].item():.2f} |'
                f"\tAux loss: {aux_loss.item():.2f}"
            )


def test_epoch(epoch, test_dataloader, model, criterion):
    model.eval()
    device = next(model.parameters()).device

    loss = AverageMeter()
    bpp_loss = AverageMeter()
    mse_loss = AverageMeter()
    aux_loss = AverageMeter()

    with torch.no_grad():
        for d in test_dataloader:
            d = d.to(device)
            out_net = model(d)
            out_criterion = criterion(out_net, d)

            aux_loss.update(model.aux_loss())
            bpp_loss.update(out_criterion["bpp_loss"])
            loss.update(out_criterion["loss"])
            mse_loss.update(out_criterion["mse_loss"])

    print(
        f"Test epoch {epoch}: Average losses:"
        f"\tLoss: {loss.avg:.3f} |"
        f"\tMSE loss: {mse_loss.avg:.3f} |"
        f"\tBpp loss: {bpp_loss.avg:.2f} |"
        f"\tAux loss: {aux_loss.avg:.2f}\n"
    )

    return loss.avg


def save_checkpoint(state, is_best, filename="checkpoint.pth.tar"):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, "checkpoint_best_loss.pth.tar")


def parse_args(argv):
    parser = argparse.ArgumentParser(description="Example training script.")
    parser.add_argument(
        "-m",
        "--model",
        default="bmshj2018-factorized",
        choices=models.keys(),
        help="Model architecture (default: %(default)s)",
    )
    parser.add_argument(
        "-d", "--dataset", type=str, required=True, help="Training dataset"
    )
    parser.add_argument(
        "-e",
        "--epochs",
        default=100,
        type=int,
        help="Number of epochs (default: %(default)s)",
    )
    parser.add_argument(
        "-lr",
        "--learning-rate",
        default=1e-4,
        type=float,
        help="Learning rate (default: %(default)s)",
    )
    parser.add_argument(
        "-n",
        "--num-workers",
        type=int,
        default=30,
        help="Dataloaders threads (default: %(default)s)",
    )
    parser.add_argument(
        "--lambda",
        dest="lmbda",
        type=float,
        default=1e-2,
        help="Bit-rate distortion parameter (default: %(default)s)",
    )
    parser.add_argument(
        "--batch-size", type=int, default=16, help="Batch size (default: %(default)s)"
    )
    parser.add_argument(
        "--test-batch-size",
        type=int,
        default=64,
        help="Test batch size (default: %(default)s)",
    )
    parser.add_argument(
        "--aux-learning-rate",
        default=1e-3,
        help="Auxiliary loss learning rate (default: %(default)s)",
    )
    parser.add_argument(
        "--patch-size",
        type=int,
        nargs=2,
        default=(256, 256),
        help="Size of the patches to be cropped (default: %(default)s)",
    )
    parser.add_argument("--cuda", action="store_true", help="Use cuda")
    parser.add_argument("--save", action="store_true", help="Save model to disk")
    parser.add_argument(
        "--seed", type=float, help="Set random seed for reproducibility"
    )
    parser.add_argument(
        "--clip_max_norm",
        default=1.0,
        type=float,
        help="gradient clipping max norm (default: %(default)s",
    )
    parser.add_argument('--checkpoint-file', type=str, help='File address to resume training from the previous saved checkpoint')
    args = parser.parse_args(argv)
    return args


def main(argv):
    args = parse_args(argv)

    if args.seed is not None:
        torch.manual_seed(args.seed)
        random.seed(args.seed)

    train_transforms = transforms.Compose(
        [transforms.RandomCrop(args.patch_size), transforms.ToTensor()]
    )

    test_transforms = transforms.Compose(
        [transforms.CenterCrop(args.patch_size), transforms.ToTensor()]
    )

    train_dataset = ImageFolder(args.dataset, split="train", transform=train_transforms)
    test_dataset = ImageFolder(args.dataset, split="test", transform=test_transforms)

    train_dataloader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        shuffle=True,
        pin_memory=True,
    )

    test_dataloader = DataLoader(
        test_dataset,
        batch_size=args.test_batch_size,
        num_workers=args.num_workers,
        shuffle=False,
        pin_memory=True,
    )

    device = "cuda" if args.cuda and torch.cuda.is_available() else "cpu"

    net = models[args.model](quality=3)
    net = net.to(device)

    if args.cuda and torch.cuda.device_count() > 1:
        net = CustomDataParallel(net)

    optimizer, aux_optimizer = configure_optimizers(net, args)
    lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min")
    criterion = RateDistortionLoss(lmbda=args.lmbda)

    last_epoch = -1
    if args.checkpoint_file:  # load from previous checkpoint
        print("Loading", args.checkpoint_file)
        checkpoint = torch.load(args.checkpoint_file, map_location=device)
        last_epoch = checkpoint["epoch"]
        net.load_state_dict((checkpoint["state_dict"]))
        net.update(force=True)  # update the model CDFs parameters.
        optimizer.load_state_dict((checkpoint["optimizer"]))
        aux_optimizer.load_state_dict((checkpoint["aux_optimizer"]))
        lr_scheduler.load_state_dict((checkpoint["lr_scheduler"]))

    best_loss = 1e10
    for epoch in range(last_epoch + 1, args.epochs):
        print(f"Learning rate: {optimizer.param_groups[0]['lr']}")
        train_one_epoch(
            net,
            criterion,
            train_dataloader,
            optimizer,
            aux_optimizer,
            epoch,
            args.clip_max_norm,
        )

        loss = test_epoch(epoch, test_dataloader, net, criterion)
        lr_scheduler.step(loss)

        is_best = loss < best_loss
        best_loss = min(loss, best_loss)
        if args.save:
            save_checkpoint(
                {
                    "epoch": epoch,
                    "state_dict": net.state_dict(),
                    "loss": loss,
                    "optimizer": optimizer.state_dict(),
                    "aux_optimizer": aux_optimizer.state_dict(),
                    "lr_scheduler": lr_scheduler.state_dict(),
                },
                is_best,
            )





if __name__ == "__main__":
    main(sys.argv[1:])

First of all, thank you for the great package!

1. Support DistributedDataParallel and DataParallel

I'm working on large-scale experiments that takes pretty long for training, and wondering if this framework can support DataParallel and DistributedDataParallel.

The current example/train.py looks like supporting Dataparallel as CustomDataParallel, but returned the following error

Traceback (most recent call last):
  File "examples/train.py", line 369, in <module>
    main(sys.argv[1:])
  File "examples/train.py", line 348, in main
    args.clip_max_norm,
  File "examples/train.py", line 159, in train_one_epoch
    out_net = model(d)
  File "/home/yoshitom/.local/share/virtualenvs/yoshitom-lJAkl1qx/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/yoshitom/.local/share/virtualenvs/yoshitom-lJAkl1qx/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 160, in forward
    replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
  File "/home/yoshitom/.local/share/virtualenvs/yoshitom-lJAkl1qx/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 165, in replicate
    return replicate(module, device_ids, not torch.is_grad_enabled())
  File "/home/yoshitom/.local/share/virtualenvs/yoshitom-lJAkl1qx/lib/python3.6/site-packages/torch/nn/parallel/replicate.py", line 140, in replicate
    param_idx = param_indices[param]
KeyError: Parameter containing:
tensor([[[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]],

        [[-10.,   0.,  10.]]], device='cuda:0', requires_grad=True)

(pipenv run python examples/train.py --data ./dataset/ --batch-size 4 --cuda on a machine with 3 GPUs)

When commenting out these two lines https://github.com/InterDigitalInc/CompressAI/blob/master/examples/train.py#L333-L334 , it looks working well

/home/yoshitom/.local/share/virtualenvs/yoshitom-lJAkl1qx/lib/python3.6/site-packages/torch/nn/modules/container.py:435: UserWarning: Setting attributes on ParameterList is not supported.
  warnings.warn("Setting attributes on ParameterList is not supported.")
Train epoch 0: [0/5000 (0%)]	Loss: 183.278 |	MSE loss: 0.278 |	Bpp loss: 2.70 |	Aux loss: 5276.71
Train epoch 0: [40/5000 (1%)]	Loss: 65.175 |	MSE loss: 0.096 |	Bpp loss: 2.70 |	Aux loss: 5273.95
Train epoch 0: [80/5000 (2%)]	Loss: 35.178 |	MSE loss: 0.050 |	Bpp loss: 2.69 |	Aux loss: 5271.21
Train epoch 0: [120/5000 (2%)]	Loss: 36.634 |	MSE loss: 0.052 |	Bpp loss: 2.68 |	Aux loss: 5268.45
Train epoch 0: [160/5000 (3%)]	Loss: 26.010 |	MSE loss: 0.036 |	Bpp loss: 2.68 |	Aux loss: 5265.67
...

Could you please fix the issue and also support DistributedDataParallel? If you need more examples to identify the components causing this issue, let me know. I have a few more examples (error messages) for both DataParallel and DistributedDataParallel with different network architectures (containing CompressionModel).

2. Publish Python package

It would be much more useful if you can publish this framework as a Python package so that we can install it with pip install compressai

Thank you!

https://github.com/InterDigitalInc/CompressAI/blob/master/compressai/models/priors.py In JointAutoregressiveHierarchicalPriors._compress_ar: ctx_p = F.conv2d( y_crop, self.context_prediction.weight, bias=self.context_prediction.bias ) self.context_prediction.weight ignores the mask, leading to the context mismatch between _compress_ar and _decompress_ar. y_hat in _compress_ar is the full feature and needs to be masked.

Solution: masked_weight = self.context_prediction.weight * self.context_prediction.mask ctx_p = F.conv2d( y_crop, masked_weight, bias=self.context_prediction.bias )

Hello,

First of all, thank you for your work.

I trained a model (bmshj2018-factorized) on my own dataset and then tried to do some tests with your evaluation script.

I used train.py, exactly as you recommend in Readme.md or in Issue 32 (now closed). python3 train.py -d $DATASET --epochs 300 -lr 1e-4 --batch-size 16 --cuda --save

Then, python3 -m compressai.utils.eval_model checkpoint /path/to/images/folder/ -a bmshj2018-factorized -p checkpoint.pth.tar

Traceback (most recent call last): File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main "main", mod_spec) File "/usr/lib/python3.7/runpy.py", line 85, in _run_code exec(code, run_globals) File "/udd/dbacchus/Codes/CompressAI/compressai/utils/eval_model/main.py", line 301, in main(sys.argv[1:]) File "/udd/dbacchus/Codes/CompressAI/compressai/utils/eval_model/main.py", line 277, in main model = load_func(*opts, run) File "/udd/dbacchus/Codes/CompressAI/compressai/utils/eval_model/main.py", line 145, in load_checkpoint return architectures[arch].from_state_dict(torch.load(checkpoint_path)).eval() File "/udd/dbacchus/Codes/CompressAI/compressai/models/priors.py", line 165, in from_state_dict N = state_dict["g_a.0.weight"].size(0) KeyError: 'g_a.0.weight'

This is the same error message as in Issue 22 but I don't understand the cause here as I'm not training a toy example not defined in the reference CompressAI models.

When I run, print(torch.load("checkpoint_best_loss.pth.tar").keys())

I get, dict_keys(['epoch', 'state_dict', 'loss', 'optimizer', 'aux_optimizer', 'lr_scheduler'])

Thank you.

Bug

Hello, I get a floating point exception when trying to update a CompressionModel. There is no stack for the error message, so I guess it comes from an internal C module.

Using prints, I could trace that the problem comes from:

model.update() -> self._pmf_to_cdf(pmf, tail_mass, pmf_length, max_length) -> _cdf = pmf_to_quantized_cdf(prob, self.entropy_coder_precision)

and that it was raised when calling the function on a Tensor p where all entries but one are zero

To Reproduce

Steps to reproduce the behavior:

Call

prob = torch.cat((p[: pmf_length[i]], tail_mass[i]), dim=0)
 _cdf = pmf_to_quantized_cdf(prob, self.entropy_coder_precision)

on the following tensor:

tensor([0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.6690e-33, 6.5444e-11, 1.0000e+00, 1.2011e-13, 1.0696e-36, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00], device='cuda:0', grad_fn=)

Expected behavior

I don't know what the returned value should be, but it seems that my problem is a corner case incorrectly handled

Environment

PyTorch version: 1.7.0 Is debug build: True CUDA used to build PyTorch: 10.2 ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.4 LTS (x86_64) GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0 Clang version: Could not collect CMake version: version 3.10.2

Python version: 3.7 (64-bit runtime) Is CUDA available: True CUDA runtime version: 9.1.85 GPU models and configuration: GPU 0: Tesla V100-SXM2-32GB GPU 1: Tesla V100-SXM2-32GB GPU 2: Tesla V100-SXM2-32GB GPU 3: Tesla V100-SXM2-32GB

Nvidia driver version: 440.33.01 cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5 HIP runtime version: N/A MIOpen runtime version: N/A

Versions of relevant libraries: [pip3] numpy==1.19.4 [pip3] pytorch-msssim==0.2.0 [pip3] torch==1.7.0 [pip3] torch-cluster==1.5.8 [pip3] torch-geometric==1.6.3 [pip3] torch-scatter==2.0.5 [pip3] torch-sparse==0.6.8 [pip3] torch-spline-conv==1.2.0 [pip3] torchvision==0.8.1 [conda] numpy 1.19.4 pypi_0 pypi [conda] pytorch-msssim 0.2.0 pypi_0 pypi [conda] torch 1.7.0 pypi_0 pypi [conda] torch-cluster 1.5.8 pypi_0 pypi [conda] torch-geometric 1.6.3 pypi_0 pypi [conda] torch-scatter 2.0.5 pypi_0 pypi [conda] torch-sparse 0.6.8 pypi_0 pypi [conda] torch-spline-conv 1.2.0 pypi_0 pypi

- PyTorch / CompressAI Version (e.g., 1.0 / 0.4.0): torch   1.7.0, compressai  1.1.5
- OS (e.g., Linux): Ubuntu 18.04.4 LTS (Bionic Beaver)
- How you installed PyTorch / CompressAI (`pip`, source): pip 
- Build command you used (if compiling from source):
- Python version: 1.7.0
- CUDA/cuDNN version: 10.2
- GPU models and configuration:
- Any other relevant information: Problem appears both on cpu and gpu

Python: 3.6.9 CUDA: 11.2 compressai: 1.1.3 torch: 1.8.1

When replicating the RD curves in README for kodak dataset (wget http://r0k.us/graphics/kodak/kodak/kodim{0,1,2}{0,1,2,3,4,5,6,7,8,9}.png in ./kodak/), I observed a weird behavior of pretrained models with CUDA.

Looks like working well without CUDA

python -m compressai.utils.eval_model pretrained ./kodak/ -a bmshj2018-hyperprior --metric mse --quality 5

metric = mse

Downloading: "https://compressai.s3.amazonaws.com/models/v1/bmshj2018-hyperprior-5-f8b614e1.pth.tar" to /home/yoshitom/.cache/torch/hub/checkpoints/bmshj2018-hyperprior-5-f8b614e1.pth.tar
100.0%
{
  "name": "bmshj2018-hyperprior",
  "description": "Inference (ans)",
  "results": {
    "psnr": [
      34.52624269077672
    ],
    "ms-ssim": [
      0.9835608204205831
    ],
    "bpp": [
      0.6686842176649305
    ],
    "encoding_time": [
      0.2404747505982717
    ],
    "decoding_time": [
      0.5095066924889883
    ]
  }
}

metric = ms-ssim

python -m compressai.utils.eval_model pretrained ./kodak/ -a bmshj2018-hyperprior --metric ms-ssim --quality 5

Downloading: "https://compressai.s3.amazonaws.com/models/v1/bmshj2018-hyperprior-ms-ssim-5-c34afc8d.pth.tar" to /home/yoshitom/.cache/torch/hub/checkpoints/bmshj2018-hyperprior-ms-ssim-5-c34afc8d.pth.tar
100.0%
{
  "name": "bmshj2018-hyperprior",
  "description": "Inference (ans)",
  "results": {
    "psnr": [
      28.992422918554542
    ],
    "ms-ssim": [
      0.9866020356615385
    ],
    "bpp": [
      0.47353786892361116
    ],
    "encoding_time": [
      0.24171670277913412
    ],
    "decoding_time": [
      0.5283569494883219
    ]
  }
}

PSNR and MS-SSIM are both NaN when using CUDA

python -m compressai.utils.eval_model pretrained ./kodak/ -a bmshj2018-hyperprior --metric mse --quality 5 --cuda

metric = mse

Downloading: "https://compressai.s3.amazonaws.com/models/v1/bmshj2018-hyperprior-5-f8b614e1.pth.tar" to /home/yoshitom/.cache/torch/hub/checkpoints/bmshj2018-hyperprior-5-f8b614e1.pth.tar
100.0%
{
  "name": "bmshj2018-hyperprior",
  "description": "Inference (ans)",
  "results": {
    "psnr": [
      NaN
    ],
    "ms-ssim": [
      NaN
    ],
    "bpp": [
      0.6686876085069443
    ],
    "encoding_time": [
      0.034142365058263145
    ],
    "decoding_time": [
      0.025616129239400227
    ]
  }
}

python -m compressai.utils.eval_model pretrained ./kodak/ -a bmshj2018-hyperprior --metric ms-ssim --quality 5 --cuda

metric = ms-ssim

Downloading: "https://compressai.s3.amazonaws.com/models/v1/bmshj2018-hyperprior-ms-ssim-5-c34afc8d.pth.tar" to /home/yoshitom/.cache/torch/hub/checkpoints/bmshj2018-hyperprior-ms-ssim-5-c34afc8d.pth.tar
100.0%
{
  "name": "bmshj2018-hyperprior",
  "description": "Inference (ans)",
  "results": {
    "psnr": [
      NaN
    ],
    "ms-ssim": [
      NaN
    ],
    "bpp": [
      0.47353786892361116
    ],
    "encoding_time": [
      0.03800355394681295
    ],
    "decoding_time": [
      0.029240707556406658
    ]
  }
}

I didn't check all the combinations (model, quality, metrics, with/without CUDA), but at least bmshj2018-hyperprior with quality=8 (besides one with quality=5) also returned NaN when using CUDA (for both mse and ms-ssim checkpoints). There may be more checkpoints that face the same issue.

When I checked the output from a model (i.e., out_dec["x_hat"]), some value in the tensor is NaN when using CUDA and that must have caused this issue.

Bug

Hello, I built a video compressor with Compressai. Based on a simple hybrid coding framework, the video compressor uses hyper-prior entropy model to compress motion and residuals separately. But when training, there will always be cases where the predicted bpp is nan randomly.

Error

Expected behavior

The hyper-prior entropy model will not predict bpp as nan during training.

Environment

- PyTorch Version: 1.7.1
- CompressAI Version: 1.1.6
- OS: Ubuntu18.04
- Python version: 3.6.2
- CUDA/cuDNN version: 11.0 / 8005

Additional context

I don't know why it will appear, and can't predict when it will appear. If I load a normal checkpoint and resume training again, it may not appear. If I finish the whole training process intermittently, the entropy model can also run normally. Could you please provide me with some help?

I try to train ScaleHyperprior model using code provide in example/train.py with Imagenet2012/DIV2K dataset. The learning rate for main optimizer and aux_optimizer both set to 1e-4. The learning rate of the main optimizer is then divided by 2 when the evaluation loss reaches a plateau as it described in https://interdigitalinc.github.io/CompressAI/zoo.html. However, atfer nearly 1 week's training, the rate-distortion performance on Kadok-24 still has a gap compared to the result provided.

I am also training model on vimeo_test_clean from Vimeo90K after 2 days, it seems will not to converge to the result provided. Have I missed something? Is there any data augmentation strategy used during training?

First of all I want to thank you for your precious work !

I have a question about aux loss. I am trying to train a video compression framework where I have a flow compressor and a residual compressor. I initialized optimizer and aux optimizer as you stated under examples directory. I set aux learning rate to 1.e-3 and it first decreases as training continues (I sum flow and residual aux losses). However at some point, it starts to increase. I tried to decrease the aux learning rate to 1.e-5 after enough iterations. Aux loss first decreases but bounces back and starts to increase. Do you have any idea what could cause this observation?

Thanks a lot

Hello,

It seems that after the recent update, I get the following error in my code

torch.nn.modules.module.ModuleAttributeError: 'FactorizedPrior' object has no attribute 'aux_parameters'

To double check, I also executed your examples/CompressAI Inference Demo.ipynb example and I got the same error. Has anything change in the structure of the code?

Thank you again for this nice library

Thanks for your great work! I trained a simple Auto Encoder with cifar10, and my loss dropped normally during training. After training, I updated the parameters of the entropy model. When I want to visualize my results, if there is no model.eval(), the reconstructed picture is normal, but if model.eval() is used, it will be very bad. Looking forward to your reply.

How to measure compressed video size?

Hi guys, I really appreciate all of your video compression work. I cannot thank you guys enough. I was wondering how I can measure the size of my compressed video. Any help or suggestions would be greatly appreciated :) Thanks.

Bug

Hi all, Thank you for a great compression library. I am using Hyperprior-architecture (Baelle ICLR 2018) but I modified the encoder and decoder (added few layers etc). I trained this modified architecture for 50 epochs on Vimeo90K triplet dataset with Quality = 2, batch_size =32. During training, the test loss was almost 0, therefore I stopped the training. However, when I evaluated the model (using the scripts available in compressAI) , the results on Kodak dataset was very unrealistic. The json is given below.

{
  "name": "bmshj2018-modified-hyperprior-mse",
  "description": "Inference (ans)",
  "results": {
    "psnr": [
      90.76645755767822
    ],
    "ms-ssim": [
      0.9999999925494194
    ],
    "bpp": [
      0.0003255208333333332
    ],
    "encoding_time": [
      0.11552197734514873
    ],
    "decoding_time": [
      0.03114013870557149
    ]
  } 

Can you please advise if this is a bug in the evaluation code or what else could be the reason for these crazy numbers? I believe i am doing something very stupid :( .

Thanks

I am trying to test kodak dataset use VTM codec with the command: python -m compressai.utils.bench vtm ./kodim -b ./VVCSoftware_VTM-VTM-9.1/bin -c ./VVCSoftware_VTM-VTM-9.1/cfg/encoder_intra_vtm.cfg -q 12 17 22 27 32 37 42 47 -j 24 >vtm9.1_kodim.json. But get different result with compressai's. The bpp is almost the same, but psnr and msssim is worse then compressai's reult. Detailed result is attached. How can I get correct result? vtm9.1_kodim.txt

Hi there, i have test cheng2020-anchor and mbt2018 model and found the encoding and decoding time cost is too high, for 256x 256 image patch, the encoding time cost is 0.8s and decoding time cost is 2.2s. My computer hardware is as follows: CPU: Intel(R) Xeon(R) CPU E5-2680 v4 @ 2.40GHz GPU: GeForce RTX 3090 Any adivce on how to speed up the encoding and decoding process?

Regarding:

https://github.com/InterDigitalInc/CompressAI/blob/1259affe66da8a05750af8ae44d90989320c2103/compressai/entropy_models/entropy_models.py#L458-L469

The following:

sign = -torch.sign(lower + upper)
sign = sign.detach()
likelihood = torch.abs(
    torch.sigmoid(sign * upper) - torch.sigmoid(sign * lower)
)

...seems to be equivalent to:

likelihood = torch.abs(torch.sigmoid(upper) - torch.sigmoid(lower))

Furthermore, since _logits_cumulative is guaranteed to be a monotonically increasing scalar function on R -> R, we must have that lower <= upper. Thus, we can remove the abs altogether:

likelihood = torch.sigmoid(upper) - torch.sigmoid(lower)

Unless I'm mistaken, there should be no difference in numerical stability either.

"Proof" by experimentation:

import torch

def likelihood_old(lower, upper):
    sign = -torch.sign(lower + upper)
    sign = sign.detach()
    likelihood = torch.abs(
        torch.sigmoid(sign * upper) - torch.sigmoid(sign * lower)
    )
    return likelihood

def likelihood_new(lower, upper):
    likelihood = torch.abs(torch.sigmoid(upper) - torch.sigmoid(lower))
    return likelihood

>>> n = 100
>>> lower, upper = (torch.randn(n) * 10**torch.randn(n), torch.randn(n) * 10**torch.randn(n))
>>> torch.isclose(f_old(lower, upper), f_new(lower, upper))
tensor([True, True, True, True, True, True, True, True, True, True, True, True,
        True, True, True, True, True, True, True, True, True, True, True, True,
        True, True, True, True, True, True, True, True, True, True, True, True,
        True, True, True, True, True, True, True, True, True, True, True, True,
        True, True, True, True, True, True, True, True, True, True, True, True,
        True, True, True, True, True, True, True, True, True, True, True, True,
        True, True, True, True, True, True, True, True, True, True, True, True,
        True, True, True, True, True, True, True, True, True, True, True, True,
        True, True, True, True])

NOTE: I haven't checked if the gradients are the same. I should probably do that...