The question about function create_compressed_model()：RuntimeError: CUDA error: device-side assert triggered

[zbnlala]

🐛 Describe the bug

Hi, I want to ask about the quantization of the non-neural network parts in the neural network.
I want to realize the quantization of the PointNet-based network to deal with classification tasks

The other part of the network are just some conv1d and relu function. Then I modified the train function with
changed dataset and evaluation method. However the error occurred . Can you help me fix this problem?
This is the error information:

Traceback (most recent call last):
  File "/home/lalala/examples/torch/classification/quant_train_cls.py", line 874, in <module>
    main(sys.argv[1:])
  File "/home/lalala/examples/torch/classification/quant_train_cls.py", line 151, in main
    start_worker(main_worker, config)
  File "/home/lalala/examples/torch/common/execution.py", line 114, in start_worker
    mp.spawn(main_worker, nprocs=config.ngpus_per_node, args=(config,))
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 241, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 197, in start_processes
    while not context.join():
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 158, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 68, in _wrap
    fn(i, *args)
  File "/home/lalala/examples/torch/classification/quant_train_cls.py", line 267, in main_worker
    compression_ctrl, model = create_compressed_model(
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/telemetry/decorator.py", line 72, in wrapped
    retval = fn(*args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/model_creation.py", line 134, in create_compressed_model
    compressed_model = builder.apply_to(nncf_network)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/compression_method_api.py", line 124, in apply_to
    transformed_model = transformer.transform(transformation_layout)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/model_transformer.py", line 78, in transform
    model.nncf.rebuild_graph()
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/nncf_network.py", line 555, in rebuild_graph
    compressed_traced_graph = builder.build_dynamic_graph(
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/graph/graph_builder.py", line 53, in build_dynamic_graph
    return tracer.trace_graph(model, context_to_use, as_eval, trace_parameters)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/dynamic_graph/graph_tracer.py", line 53, in trace_graph
    self.custom_forward_fn(model)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/dynamic_graph/graph_tracer.py", line 96, in default_dummy_forward_fn
    retval = model(*args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/nncf_network.py", line 1004, in __call__
    return ORIGINAL_CALL(self, *args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/nncf_network.py", line 1036, in forward
    retval = wrap_module_call(self.nncf._original_unbound_forward)(self, *args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/dynamic_graph/wrappers.py", line 154, in wrapped
    retval = module_call(self, *args, **kwargs)
  File "/home/lalala/examples/torch/classification/models/pointnet2_cls.py", line 39, in forward
    l1_xyz, l1_points = self.sa1(xyz, norm)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/dynamic_graph/wrappers.py", line 154, in wrapped
    retval = module_call(self, *args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/lalala/examples/torch/classification/models/pointnet2_utils.py", line 197, in forward
    new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
  File "/home/lalala/examples/torch/classification/models/pointnet2_utils.py", line 131, in sample_and_group
    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
  File "/home/lalala/examples/torch/classification/models/pointnet2_utils.py", line 59, in index_points
    new_points = points[batch_indices, idx, :]
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/dynamic_graph/wrappers.py", line 98, in wrapped
    result = _execute_op(op_address, operator_info, operator, ctx, *args, **kwargs)
  File "/home/lalala/anaconda3/envs/quantization/lib/python3.9/site-packages/nncf/torch/dynamic_graph/wrappers.py", line 179, in _execute_op
    result = operator(*args, **kwargs)
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Environment

nccf==2.10.0
pytorch==2.2.2
python==3.9

Minimal Reproducible Example

The non-neural network parts of the network is as follow:

def square_distance(src, dst):
    Calculate Euclid distance between each two points.
    Input:
        src: source points, [B, N, C]
        dst: target points, [B, M, C]
    Output:
        dist: per-point square distance, [B, N, M]

    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dis

 
def index_points(points, idx):
    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]

    device = points.device
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points

def farthest_point_sample(xyz, npoint):

    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]

    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10

    batch_indices = torch.arange(B, dtype=torch.long).to(device)
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    return centroids


def query_ball_point(radius, nsample, xyz, new_xyz):

    Input:
        radius: local region radius
        nsample: max sample number in local region
        xyz: all points, [B, N, 3]
        new_xyz: query points, [B, S, 3]
    Return:
        group_idx: grouped points index, [B, S, nsample]

    device = xyz.device
    B, N, C = xyz.shape
    _, S, _ = new_xyz.shape
    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
    sqrdists = square_distance(new_xyz, xyz)
    group_idx[sqrdists > radius ** 2] = N
    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
    mask = group_idx == N
    group_idx[mask] = group_first[mask]
    return group_idx


def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False):
    Input:
        npoint:
        radius:
        nsample:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, npoint, nsample, 3]
        new_points: sampled points data, [B, npoint, nsample, 3+D]

    B, N, C = xyz.shape
    S = npoint
    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
    new_xyz = index_points(xyz, fps_idx) # [B, npoint, 3]
    idx = query_ball_point(radius, nsample, xyz, new_xyz)
    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)


    if points is not None:
        grouped_points = index_points(points, idx)
        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
    else:
        new_points = grouped_xyz_norm
    if returnfps:
        return new_xyz, new_points, grouped_xyz, fps_idx
    else:
        return new_xyz, new_points


def sample_and_group_all(xyz, points):

    Input:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, 1, 3]
        new_points: sampled points data, [B, 1, N, 3+D]

    device = xyz.device
    B, N, C = xyz.shape
    new_xyz = torch.zeros(B, 1, C).to(device)
    grouped_xyz = xyz.view(B, 1, N, C)
    if points is not None:
        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
    else:
        new_points = grouped_xyz
    return new_xyz, new_points


class PointNetSetAbstraction(nn.Module):
    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
        super(PointNetSetAbstraction, self).__init__()
        self.npoint = npoint
        self.radius = radius
        self.nsample = nsample
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
            last_channel = out_channel
        self.group_all = group_all

    def forward(self, xyz, points):

        Input:
            xyz: input points position data, [B, C, N]
            points: input points data, [B, D, N]
        Return:
            new_xyz: sampled points position data, [B, C, S]
            new_points_concat: sample points feature data, [B, D', S]

        xyz = xyz.permute(0, 2, 1)
        if points is not None:
            points = points.permute(0, 2, 1)

        if self.group_all:
            new_xyz, new_points = sample_and_group_all(xyz, points)
        else:
            new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
        # new_xyz: sampled points position data, [B, npoint, C]
        # new_points: sampled points data, [B, npoint, nsample, C+D]
        # print(new_points.shape)

        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points =  F.relu(bn(conv(new_points)))

        new_points = torch.max(new_points, 2)[0] # max at dimension of nsample 
        new_xyz = new_xyz.permute(0, 2, 1)
        return new_xyz, new_points

[AlexanderDokuchaev]

HI @zbnlala
Could you please provide a full script to reproduce the issue?

[zbnlala]

HI @AlexanderDokuchaev. Sure!
First, my code is modified from examples/torch/classification and put the train.py into this dir.
Note that, for train.py you just need to focus on the function main_worker().
Moreover, the modified dataset code is ModelNetDataLoader.py placed in examples/torch/classification/data_utils(the dir is created by myself). From now the code must run on the modelnet dataset, I do not know how to run create_compressed_model() without the dataloader. Therefore how to get dataset is here.

the model architecture is described in pointnet2_cls.py and pointnet2_utils.py placed in
examples/torch/classification/models.
And then, the json is pointnet_v2_classification_int8.json placed in examples/torch/classification/configs/quantization.

Finally , run the script will cause the bug
CUDA_VISIBLE_DEVICES="0,1" NCCL_P2P_DISABLE=1 python train.py -m test --config configs/quantization/pointnet_v2_classification_int8.json --log-dir=../../results/quantization/pointnet_v2_int8/ --multiprocessing-distributed

However, if I run script with extra --cpu-only, there will no bug.(CUDA version 12.3)
train.py

# Copyright (c) 2024 Intel Corporation
# 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 os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import os.path as osp
import sys
import time
import warnings
from copy import deepcopy
from functools import partial
from pathlib import Path
from shutil import copyfile
from typing import Any
from tqdm import tqdm
import numpy as np

import torch
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
from torch import nn
from torch.backends import cudnn
from torch.cuda.amp.autocast_mode import autocast
from torch.nn.modules.loss import _Loss
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torchvision import datasets
from torchvision import models
from torchvision import transforms
from torchvision.datasets import CIFAR10
from torchvision.datasets import CIFAR100
from torchvision.models import InceptionOutputs
from examples.common.paths import configure_paths
from examples.common.sample_config import SampleConfig
from examples.common.sample_config import create_sample_config
from examples.torch.common.argparser import get_common_argument_parser
from examples.torch.common.argparser import parse_args
from examples.torch.common.example_logger import logger
from examples.torch.common.execution import ExecutionMode
from examples.torch.common.execution import get_execution_mode
from examples.torch.common.execution import prepare_model_for_execution
from examples.torch.common.execution import set_seed
from examples.torch.common.execution import start_worker
from examples.torch.common.export import export_model
from examples.torch.common.model_loader import COMPRESSION_STATE_ATTR
from examples.torch.common.model_loader import MODEL_STATE_ATTR
from examples.torch.common.model_loader import extract_model_and_compression_states
from examples.torch.common.model_loader import load_model
from examples.torch.common.model_loader import load_resuming_checkpoint
from examples.torch.common.optimizer import get_parameter_groups
from examples.torch.common.optimizer import make_optimizer
from examples.torch.common.utils import MockDataset
from examples.torch.common.utils import NullContextManager
from examples.torch.common.utils import configure_device
from examples.torch.common.utils import configure_logging
from examples.torch.common.utils import create_code_snapshot
from examples.torch.common.utils import get_run_name
from examples.torch.common.utils import is_pretrained_model_requested
from examples.torch.common.utils import is_staged_quantization
from examples.torch.common.utils import make_additional_checkpoints
from examples.torch.common.utils import print_args
from examples.torch.common.utils import write_metrics
from nncf.api.compression import CompressionStage
from nncf.common.accuracy_aware_training import create_accuracy_aware_training_loop
from nncf.common.utils.tensorboard import prepare_for_tensorboard
from nncf.config.utils import is_accuracy_aware_training
from nncf.torch import create_compressed_model
from nncf.torch.checkpoint_loading import load_state
from nncf.torch.dynamic_graph.io_handling import FillerInputInfo
from nncf.torch.initialization import default_criterion_fn
from nncf.torch.initialization import register_default_init_args
from nncf.torch.structures import ExecutionParameters
from nncf.torch.utils import is_main_process
from nncf.torch.utils import safe_thread_call

from examples.torch.classification.models.pointnet2_cls import pointnetv2_cls
from examples.torch.common import restricted_pickle_module

from data_utils.ModelNetDataLoader import ModelNetDataLoader

model_names = sorted(
    name
    for name, val in models.__dict__.items()
    if name.islower() and not name.startswith("__") and callable(val)
)


def get_argument_parser():
    parser = get_common_argument_parser()
    parser.add_argument(
        "--dataset",
        help="Dataset to use.",
        choices=["imagenet", "cifar100", "cifar10"],
        default=None,
    )
    parser.add_argument(
        "--local_rank",
        default=None,
    )
    parser.add_argument(
        "--test-every-n-epochs",
        default=1,
        type=int,
        help="Enables running validation every given number of epochs",
    )
    parser.add_argument(
        "--mixed-precision",
        dest="mixed_precision",
        help="Enables torch.cuda.amp autocasting during training and validation steps",
        action="store_true",
    )
    return parser


def main(argv):
    parser = get_argument_parser()
    args = parse_args(parser, argv)
    config = create_sample_config(args, parser)

    if config.dist_url == "env://":
        config.update_from_env()

    configure_paths(config, get_run_name(config))
    copyfile(args.config, osp.join(config.log_dir, "config.json"))
    source_root = Path(__file__).absolute().parents[2]  # nncf root
    create_code_snapshot(source_root, osp.join(config.log_dir, "snapshot.tar.gz"))

    if config.seed is not None:
        warnings.warn(
            "You have chosen to seed training. "
            "This will turn on the CUDNN deterministic setting, "
            "which can slow down your training considerably! "
            "You may see unexpected behavior when restarting "
            "from checkpoints."
        )

    config.execution_mode = get_execution_mode(config)

    if config.metrics_dump is not None:
        write_metrics(0, config.metrics_dump)

    if not is_staged_quantization(config):
        start_worker(main_worker, config)
    else:
        from examples.torch.classification.staged_quantization_worker import (
            staged_quantization_main_worker,
        )

        start_worker(staged_quantization_main_worker, config)


def inception_criterion_fn(
    model_outputs: Any, target: Any, criterion: _Loss
) -> torch.Tensor:
    # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958
    output, aux_outputs = model_outputs
    loss1 = criterion(output, target)
    loss2 = criterion(aux_outputs, target)
    return loss1 + 0.4 * loss2


def main_worker(current_gpu, config: SampleConfig):
    configure_device(current_gpu, config)
    if is_main_process():
        configure_logging(logger, config)
        print_args(config)
    else:
        config.tb = None

    set_seed(config)

    # define loss function (criterion)
    criterion = nn.CrossEntropyLoss()
    criterion = criterion.to(config.device)

    model_name = config["model"]
    train_criterion_fn = (
        inception_criterion_fn if "inception" in model_name else default_criterion_fn
    )

    train_loader = train_sampler = val_loader = None
    resuming_checkpoint_path = config.resuming_checkpoint_path
    nncf_config = config.nncf_config
    pretrained = is_pretrained_model_requested(config)
    is_export_only = "export" in config.mode and (
        "train" not in config.mode and "test" not in config.mode
    )

    if is_export_only:
        assert pretrained or (resuming_checkpoint_path is not None)
    else:
        # Data loading code
        # train_dataset, val_dataset = create_datasets(config)
        train_dataset, val_dataset = pointnet_dataset(config)
        train_loader, train_sampler, val_loader, init_loader = create_data_loaders(
            config, train_dataset, val_dataset
        )

        def train_steps_fn(loader, model, optimizer, compression_ctrl, train_steps):
            train_epoch(
                loader,
                model,
                criterion,
                train_criterion_fn,
                optimizer,
                compression_ctrl,
                0,
                config,
                train_iters=train_steps,
                log_training_info=False,
            )

        def validate_model_fn(model, eval_loader):
            instance_acc,class_acc  = validate(
                eval_loader, model, criterion, config
            )
            return instance_acc,class_acc

        def model_eval_fn(model):
            acc,_ = validate(val_loader, model, criterion, config)
            return acc

        execution_params = ExecutionParameters(config.cpu_only, config.current_gpu)

        nncf_config = register_default_init_args(
            nncf_config,
            init_loader,
            criterion=criterion,
            criterion_fn=train_criterion_fn,
            train_steps_fn=train_steps_fn,
            validate_fn=lambda *x: validate_model_fn(*x)[::2],
            autoq_eval_fn=lambda *x: validate_model_fn(*x)[1],
            val_loader=val_loader,
            model_eval_fn=model_eval_fn,
            device=config.device,
            execution_parameters=execution_params,
        )

    # create model
    num_classes=config.get("num_classes", 1000)
    model_params=config.get("model_params")
    weights_path=config.get("weights")
    load_model_fn = partial(pointnetv2_cls, num_class=num_classes,pretrained=pretrained,load_path=weights_path)
    model = safe_thread_call(load_model_fn)
    if not pretrained and weights_path is not None:
        # Check if provided path is a url and download the checkpoint if yes
        sd = torch.load(weights_path, map_location="cpu", pickle_module=restricted_pickle_module)
       
        sd=sd["model_state_dict"]
        if MODEL_STATE_ATTR in sd:
            sd = sd[MODEL_STATE_ATTR]
        load_state(model, sd, is_resume=False)

    model.to(config.device)

    if "train" in config.mode and is_accuracy_aware_training(config):
        uncompressed_model_accuracy = model_eval_fn(model)

    resuming_checkpoint = None
    if resuming_checkpoint_path is not None:
        resuming_checkpoint = load_resuming_checkpoint(resuming_checkpoint_path)
    model_state_dict, compression_state = extract_model_and_compression_states(
        resuming_checkpoint
    )
    compression_ctrl, model = create_compressed_model(
        model, nncf_config, compression_state
    )


def train(
    config,
    compression_ctrl,
    model,
    criterion,
    criterion_fn,
    lr_scheduler,
    model_name,
    optimizer,
    train_loader,
    train_sampler,
    val_loader,
    best_acc1=0,
):
    best_compression_stage = CompressionStage.UNCOMPRESSED

    for epoch in range(config.start_epoch, config.epochs):
        # update compression scheduler state at the begin of the epoch
        compression_ctrl.scheduler.epoch_step()

        if config.distributed:
            train_sampler.set_epoch(epoch)

        # train for one epoch
        train_epoch(
            train_loader,
            model,
            criterion,
            criterion_fn,
            optimizer,
            compression_ctrl,
            epoch,
            config,
        )

        # Learning rate scheduling should be applied after optimizer’s update
        lr_scheduler.step(
            epoch if not isinstance(lr_scheduler, ReduceLROnPlateau) else best_acc1
        )

        # compute compression algo statistics
        statistics = compression_ctrl.statistics()

        acc1 = best_acc1

        best_instance_acc = 0.0
        best_class_acc = 0.0
        if epoch % config.test_every_n_epochs == 0:
            # evaluate on validation set
            instance_acc, class_acc = validate(val_loader, model, criterion, config, epoch=epoch)
        if (instance_acc >= best_instance_acc):
                best_instance_acc = instance_acc
                best_epoch = epoch + 1

        if (class_acc >= best_class_acc):
            best_class_acc = class_acc
        logger.info('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc))
        logger.info('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc))
        compression_stage = compression_ctrl.compression_stage()
        # remember best acc@1, considering compression stage. If current acc@1 less then the best acc@1, checkpoint
        # still can be best if current compression stage is larger than the best one. Compression stages in ascending
        # order: UNCOMPRESSED, PARTIALLY_COMPRESSED, FULLY_COMPRESSED.
        is_best_by_accuracy = (
            acc1 > best_acc1 and compression_stage == best_compression_stage
        )
        is_best = is_best_by_accuracy or compression_stage > best_compression_stage
        if is_best:
            best_acc1 = acc1
        best_compression_stage = max(compression_stage, best_compression_stage)
        if is_main_process():
            logger.info(statistics.to_str())

            if config.metrics_dump is not None:
                acc = best_acc1 / 100
                write_metrics(acc, config.metrics_dump)

            checkpoint_path = osp.join(
                config.checkpoint_save_dir, get_run_name(config) + "_last.pth"
            )
            checkpoint = {
                "epoch": epoch + 1,
                "arch": model_name,
                MODEL_STATE_ATTR: model.state_dict(),
                COMPRESSION_STATE_ATTR: compression_ctrl.get_compression_state(),
                "best_acc1": best_acc1,
                "acc1": acc1,
                "optimizer": optimizer.state_dict(),
            }

            torch.save(checkpoint, checkpoint_path)
            make_additional_checkpoints(checkpoint_path, is_best, epoch + 1, config)

            for key, value in prepare_for_tensorboard(statistics).items():
                config.tb.add_scalar(
                    "compression/statistics/{0}".format(key),
                    value,
                    len(train_loader) * epoch,
                )


def get_dataset(dataset_config, config, transform, is_train):
    if dataset_config == "imagenet":
        prefix = "train" if is_train else "val"
        return datasets.ImageFolder(osp.join(config.dataset_dir, prefix), transform)
    # For testing purposes
    num_images = config.get("num_mock_images", 1000)
    if dataset_config == "mock_32x32":
        return MockDataset(
            img_size=(32, 32), transform=transform, num_images=num_images
        )
    if dataset_config == "mock_299x299":
        return MockDataset(
            img_size=(299, 299), transform=transform, num_images=num_images
        )
    return create_cifar(config, dataset_config, is_train, transform)


def create_cifar(config, dataset_config, is_train, transform):
    create_cifar_fn = None
    if dataset_config in ["cifar100", "cifar100_224x224"]:
        create_cifar_fn = partial(
            CIFAR100, config.dataset_dir, train=is_train, transform=transform
        )
    if dataset_config == "cifar10":
        create_cifar_fn = partial(
            CIFAR10, config.dataset_dir, train=is_train, transform=transform
        )
    if create_cifar_fn:
        return safe_thread_call(
            partial(create_cifar_fn, download=True),
            partial(create_cifar_fn, download=False),
        )
    return None


def create_datasets(config):
    dataset_config = config.dataset if config.dataset is not None else "imagenet"
    dataset_config = dataset_config.lower()
    assert dataset_config in [
        "imagenet",
        "cifar100",
        "cifar10",
        "cifar100_224x224",
        "mock_32x32",
        "mock_299x299",
    ], "Unknown dataset option"

    if dataset_config == "imagenet":
        normalize = transforms.Normalize(
            mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)
        )
    elif dataset_config in ["cifar100", "cifar100_224x224"]:
        normalize = transforms.Normalize(
            mean=(0.5071, 0.4865, 0.4409), std=(0.2673, 0.2564, 0.2761)
        )
    elif dataset_config == "cifar10":
        normalize = transforms.Normalize(
            mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616)
        )
    elif dataset_config in ["mock_32x32", "mock_299x299"]:
        normalize = transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))

    input_info = FillerInputInfo.from_nncf_config(config)
    image_size = input_info.elements[0].shape[-1]
    size = int(image_size / 0.875)
    if dataset_config in ["cifar10", "cifar100_224x224", "cifar100"]:
        list_val_transforms = [transforms.ToTensor(), normalize]
        if dataset_config == "cifar100_224x224":
            list_val_transforms.insert(0, transforms.Resize(image_size))
        val_transform = transforms.Compose(list_val_transforms)

        list_train_transforms = [
            transforms.RandomCrop(image_size, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
        if dataset_config == "cifar100_224x224":
            list_train_transforms.insert(0, transforms.Resize(image_size))
        train_transforms = transforms.Compose(list_train_transforms)
    elif dataset_config in ["mock_32x32", "mock_299x299"]:
        val_transform = transforms.Compose(
            [
                transforms.Resize(size),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                normalize,
            ]
        )
        train_transforms = transforms.Compose(
            [
                transforms.Resize(size),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                normalize,
            ]
        )
    else:
        val_transform = transforms.Compose(
            [
                transforms.Resize(size),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                normalize,
            ]
        )
        train_transforms = transforms.Compose(
            [
                transforms.RandomResizedCrop(image_size),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]
        )

    val_dataset = get_dataset(dataset_config, config, val_transform, is_train=False)
    train_dataset = get_dataset(dataset_config, config, train_transforms, is_train=True)

    return train_dataset, val_dataset

def pointnet_dataset(config):
    data_path=config.get("dataset_dir")
    num_point =  config.get("num_point")
    use_uniform_sample =  config.get("use_uniform_sample")
    use_normals =  config.get("use_normals")
    num_category =  config.get("num_classes")
    process_data=config.get("process_data")
    train_dataset = ModelNetDataLoader(root=data_path, args=None, split='train', 
                                       process_data=process_data,num_point=num_point,
                                       use_uniform_sample=use_uniform_sample,
                                       use_normals=use_normals,num_category=num_category)
    val_dataset = ModelNetDataLoader(root=data_path, args=None, split='test', 
                                     process_data=process_data,num_point=num_point,
                                     use_uniform_sample=use_uniform_sample,
                                     use_normals=use_normals,num_category=num_category)
    return train_dataset, val_dataset

def create_data_loaders(config, train_dataset, val_dataset):
    pin_memory = config.execution_mode != ExecutionMode.CPU_ONLY
    # When using a single GPU per process and per
    # DistributedDataParallel, we need to divide the batch size
    # ourselves based on the total number of GPUs we have
    batch_size = int(config.batch_size)
    workers = int(config.workers)
    batch_size_val = (
        int(config.batch_size_val)
        if config.batch_size_val is not None
        else int(config.batch_size)
    )
    if config.execution_mode == ExecutionMode.MULTIPROCESSING_DISTRIBUTED:
        batch_size //= config.ngpus_per_node
        batch_size_val //= config.ngpus_per_node
        workers //= config.ngpus_per_node

    val_sampler = torch.utils.data.SequentialSampler(val_dataset)
    val_loader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=batch_size_val,
        shuffle=False,
        num_workers=workers,
        pin_memory=pin_memory,
        sampler=val_sampler,
        drop_last=False,
    )

    train_sampler = None
    if config.distributed:
        sampler_seed = 0 if config.seed is None else config.seed
        dist_sampler_shuffle = config.seed is None
        train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_dataset, seed=sampler_seed, shuffle=dist_sampler_shuffle
        )

    train_shuffle = train_sampler is None and config.seed is None

    def create_train_data_loader(batch_size_):
        return torch.utils.data.DataLoader(
            train_dataset,
            batch_size=batch_size_,
            shuffle=train_shuffle,
            num_workers=workers,
            pin_memory=pin_memory,
            sampler=train_sampler,
            drop_last=True,
        )

    train_loader = create_train_data_loader(batch_size)

    if config.batch_size_init:
        init_loader = create_train_data_loader(config.batch_size_init)
    else:
        init_loader = deepcopy(train_loader)
    return train_loader, train_sampler, val_loader, init_loader


def train_epoch(
    train_loader,
    model,
    criterion,
    criterion_fn,
    optimizer,
    compression_ctrl,
    epoch,
    config,
    train_iters=None,
    log_training_info=True,
):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    compression_losses = AverageMeter()
    criterion_losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    if train_iters is None:
        train_iters = len(train_loader)

    compression_scheduler = compression_ctrl.scheduler
    casting = autocast if config.mixed_precision else NullContextManager

    # switch to train mode
    model.train()

    end = time.time()
    for i, (input_, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        compression_scheduler.step()

        input_ = input_.to(config.device)
        target = target.to(config.device)

        # compute output
        with casting():
            output = model(input_)
            criterion_loss = criterion_fn(output, target, criterion)

            # compute compression loss
            compression_loss = compression_ctrl.loss()
            loss = criterion_loss + compression_loss

        if isinstance(output, InceptionOutputs):
            output = output.logits
        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        losses.update(loss.item(), input_.size(0))
        comp_loss_val = (
            compression_loss.item()
            if isinstance(compression_loss, torch.Tensor)
            else compression_loss
        )
        compression_losses.update(comp_loss_val, input_.size(0))
        criterion_losses.update(criterion_loss.item(), input_.size(0))
        top1.update(acc1, input_.size(0))
        top5.update(acc5, input_.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % config.print_freq == 0 and log_training_info:
            logger.info(
                "{rank}: "
                "Epoch: [{0}][{1}/{2}] "
                "Lr: {3:.3} "
                "Time: {batch_time.val:.3f} ({batch_time.avg:.3f}) "
                "Data: {data_time.val:.3f} ({data_time.avg:.3f}) "
                "CE_loss: {ce_loss.val:.4f} ({ce_loss.avg:.4f}) "
                "CR_loss: {cr_loss.val:.4f} ({cr_loss.avg:.4f}) "
                "Loss: {loss.val:.4f} ({loss.avg:.4f}) "
                "Acc@1: {top1.val:.3f} ({top1.avg:.3f}) "
                "Acc@5: {top5.val:.3f} ({top5.avg:.3f})".format(
                    epoch,
                    i,
                    len(train_loader),
                    get_lr(optimizer),
                    batch_time=batch_time,
                    data_time=data_time,
                    ce_loss=criterion_losses,
                    cr_loss=compression_losses,
                    loss=losses,
                    top1=top1,
                    top5=top5,
                    rank=(
                        "{}:".format(config.rank)
                        if config.multiprocessing_distributed
                        else ""
                    ),
                )
            )

        if is_main_process() and log_training_info:
            global_step = train_iters * epoch
            config.tb.add_scalar(
                "train/learning_rate", get_lr(optimizer), i + global_step
            )
            config.tb.add_scalar(
                "train/criterion_loss", criterion_losses.val, i + global_step
            )
            config.tb.add_scalar(
                "train/compression_loss", compression_losses.val, i + global_step
            )
            config.tb.add_scalar("train/loss", losses.val, i + global_step)
            config.tb.add_scalar("train/top1", top1.val, i + global_step)
            config.tb.add_scalar("train/top5", top5.val, i + global_step)

            statistics = compression_ctrl.statistics(quickly_collected_only=True)
            for stat_name, stat_value in prepare_for_tensorboard(statistics).items():
                config.tb.add_scalar(
                    "train/statistics/{}".format(stat_name), stat_value, i + global_step
                )

        if i >= train_iters:
            break


def validate(val_loader, model, criterion, config, epoch=0, log_validation_info=True):
    mean_correct = []
    num_classes=config.get("num_classes")
    class_acc = np.zeros((num_classes, 3))
    # switch to evaluate mode
    model.eval()

    casting = autocast if config.mixed_precision else NullContextManager
    with torch.no_grad():
        
        for j, (points, target) in tqdm(enumerate(val_loader),total=len(val_loader)):
            points, target = points.to(config.device), target.to(config.device)
            with casting():
                pred, _ = model(points)
            # print(pred)
            pred_choice = pred.data.max(1)[1]

            for cat in np.unique(target.cpu()):
                classacc = pred_choice[target == cat].eq(target[target == cat].long().data).cpu().sum()
                class_acc[cat, 0] += classacc.item() / float(points[target == cat].size()[0])
                class_acc[cat, 1] += 1

            correct = pred_choice.eq(target.long().data).cpu().sum()
            mean_correct.append(correct.item() / float(points.size()[0]))
    class_acc[:, 2] = class_acc[:, 0] / class_acc[:, 1]
    class_acc = np.mean(class_acc[:, 2])
    instance_acc = np.mean(mean_correct)
    return instance_acc,class_acc


class AverageMeter:
    """Computes and stores the average and current value"""

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

    def reset(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


def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].reshape(-1).sum(0, keepdim=True)
            res.append(correct_k.float().mul_(100.0 / batch_size).item())
        return res


def get_lr(optimizer):
    return optimizer.param_groups[0]["lr"]


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

ModelNetDataLoader.py

import os
import numpy as np
import warnings
import pickle

from tqdm import tqdm
from torch.utils.data import Dataset

warnings.filterwarnings('ignore')


def pc_normalize(pc):
    centroid = np.mean(pc, axis=0)
    pc = pc - centroid
    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
    pc = pc / m
    return pc


def farthest_point_sample(point, npoint):
    """
    Input:
        xyz: pointcloud data, [N, D]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [npoint, D]
    """
    N, D = point.shape
    xyz = point[:,:3]
    centroids = np.zeros((npoint,))
    distance = np.ones((N,)) * 1e10
    farthest = np.random.randint(0, N)
    for i in range(npoint):
        centroids[i] = farthest
        centroid = xyz[farthest, :]
        dist = np.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = np.argmax(distance, -1)
    point = point[centroids.astype(np.int32)]
    return point


class ModelNetDataLoader(Dataset):
    def __init__(self, root, args, split='train', process_data=False,num_point=1024,use_uniform_sample=None,use_normals=False,num_category=40):
        self.root = root
        if args is not None:
            self.npoints = args.num_point
            self.uniform = args.use_uniform_sample
            self.use_normals = args.use_normals
            self.num_category = args.num_category
        else:
            self.npoints = num_point
            self.uniform = use_uniform_sample
            self.use_normals = use_normals
            self.num_category = num_category
        self.process_data = process_data

        if self.num_category == 10:
            self.catfile = os.path.join(self.root, 'modelnet10_shape_names.txt')
        else:
            self.catfile = os.path.join(self.root, 'modelnet40_shape_names.txt')

        self.cat = [line.rstrip() for line in open(self.catfile)]
        self.classes = dict(zip(self.cat, range(len(self.cat))))

        shape_ids = {}
        if self.num_category == 10:
            shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_train.txt'))]
            shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_test.txt'))]
        else:
            shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_train.txt'))]
            shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_test.txt'))]

        assert (split == 'train' or split == 'test')
        shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]]
        self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], shape_ids[split][i]) + '.txt') for i
                         in range(len(shape_ids[split]))]
        print('The size of %s data is %d' % (split, len(self.datapath)))

        if self.uniform:
            self.save_path = os.path.join(root, 'modelnet%d_%s_%dpts_fps.dat' % (self.num_category, split, self.npoints))
        else:
            self.save_path = os.path.join(root, 'modelnet%d_%s_%dpts.dat' % (self.num_category, split, self.npoints))

        if self.process_data:
            if not os.path.exists(self.save_path):
                print('Processing data %s (only running in the first time)...' % self.save_path)
                self.list_of_points = [None] * len(self.datapath)
                self.list_of_labels = [None] * len(self.datapath)

                for index in tqdm(range(len(self.datapath)), total=len(self.datapath)):
                    fn = self.datapath[index]
                    cls = self.classes[self.datapath[index][0]]
                    cls = np.array([cls]).astype(np.int32)
                    point_set = np.loadtxt(fn[1], delimiter=',').astype(np.float32)

                    if self.uniform:
                        point_set = farthest_point_sample(point_set, self.npoints)
                    else:
                        point_set = point_set[0:self.npoints, :]

                    self.list_of_points[index] = point_set
                    self.list_of_labels[index] = cls

                with open(self.save_path, 'wb') as f:
                    pickle.dump([self.list_of_points, self.list_of_labels], f)
            else:
                print('Load processed data from %s...' % self.save_path)
                with open(self.save_path, 'rb') as f:
                    self.list_of_points, self.list_of_labels = pickle.load(f)

    def __len__(self):
        return len(self.datapath)

    def _get_item(self, index):
        if self.process_data:
            point_set, label = self.list_of_points[index], self.list_of_labels[index]
        else:
            fn = self.datapath[index]
            cls = self.classes[self.datapath[index][0]]
            label = np.array([cls]).astype(np.int32)
            point_set = np.loadtxt(fn[1], delimiter=',').astype(np.float32)

            if self.uniform:
                point_set = farthest_point_sample(point_set, self.npoints)
            else:
                point_set = point_set[0:self.npoints, :]
                
        point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
        if not self.use_normals:
            point_set = point_set[:, 0:3]

        return point_set, label[0]

    def __getitem__(self, index):
        return self._get_item(index)


if __name__ == '__main__':
    import torch

    data = ModelNetDataLoader('/data/modelnet40_normal_resampled/', split='train')
    DataLoader = torch.utils.data.DataLoader(data, batch_size=12, shuffle=True)
    for point, label in DataLoader:
        print(point.shape)
        print(label.shape)

pointnet2_utils.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from time import time
import numpy as np

def timeit(tag, t):
    print("{}: {}s".format(tag, time() - t))
    return time()

def pc_normalize(pc):
    l = pc.shape[0]
    centroid = np.mean(pc, axis=0)
    pc = pc - centroid
    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
    pc = pc / m
    return pc

def square_distance(src, dst):
    """
    Calculate Euclid distance between each two points.

    src^T * dst = xn * xm + yn * ym + zn * zm；
    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst

    Input:
        src: source points, [B, N, C]
        dst: target points, [B, M, C]
    Output:
        dist: per-point square distance, [B, N, M]
    """
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist


def index_points(points, idx):
    """

    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]
    """
    device = points.device
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points

#1
def farthest_point_sample(xyz, npoint):
    """
    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    
    # Modified: take the farthest point to (0, 0, 0) as the first sample rather than random sampling
    # farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
    farthest = torch.argmax(torch.square(xyz).sum(dim=2), dim=1)    

    batch_indices = torch.arange(B, dtype=torch.long).to(device)
    for i in range(npoint):
        centroids[:, i] = farthest
        cent = xyz[batch_indices, farthest, :]
        if cent.shape[-1]!=3:
            print("cent.shape",cent.shape)
            print("xyz.shape",xyz.shape)
        cent = cent.view(B, 1, 3)
        dist = torch.sum((xyz - cent) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    return centroids

#2
def query_ball_point(radius, nsample, xyz, new_xyz):
    """
    Input:
        radius: local region radius
        nsample: max sample number in local region
        xyz: all points, [B, N, 3]
        new_xyz: query points, [B, S, 3]
    Return:
        group_idx: grouped points index, [B, S, nsample]
    """
    device = xyz.device
    B, N, C = xyz.shape
    _, S, _ = new_xyz.shape
    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
    sqrdists = square_distance(new_xyz, xyz)
    group_idx[sqrdists > radius ** 2] = N
    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
    mask = group_idx == N
    group_idx[mask] = group_first[mask]
    return group_idx

#3
def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False):
    """
    Input:
        npoint:
        radius:
        nsample:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, npoint, nsample, 3]
        new_points: sampled points data, [B, npoint, nsample, 3+D]
    """
    B, N, C = xyz.shape
    S = npoint
    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
    new_xyz = index_points(xyz, fps_idx) # [B, npoint, 3]
    idx = query_ball_point(radius, nsample, xyz, new_xyz)
    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)


    if points is not None:
        grouped_points = index_points(points, idx)
        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
    else:
        new_points = grouped_xyz_norm
    if returnfps:
        return new_xyz, new_points, grouped_xyz, fps_idx
    else:
        return new_xyz, new_points

# simply concatnate the xyz coordinate and point feature
def sample_and_group_all(xyz, points):
    """
    Input:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, 1, 3]
        new_points: sampled points data, [B, 1, N, 3+D]
    """
    device = xyz.device
    B, N, C = xyz.shape
    new_xyz = torch.zeros(B, 1, C).to(device)
    grouped_xyz = xyz.view(B, 1, N, C)
    if points is not None:
        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
    else:
        new_points = grouped_xyz
    return new_xyz, new_points

#4
class PointNetSetAbstraction(nn.Module):
    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
        super(PointNetSetAbstraction, self).__init__()
        self.npoint = npoint
        self.radius = radius
        self.nsample = nsample
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
            last_channel = out_channel
        self.group_all = group_all

    def forward(self, xyz, points):
        """
        Input:
            xyz: input points position data, [B, C, N]
            points: input points data, [B, D, N]
        Return:
            new_xyz: sampled points position data, [B, C, S]
            new_points_concat: sample points feature data, [B, D', S]
        """
        xyz = xyz.permute(0, 2, 1)
        if points is not None:
            points = points.permute(0, 2, 1)

        if self.group_all:
            new_xyz, new_points = sample_and_group_all(xyz, points)
        else:
            new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
        # new_xyz: sampled points position data, [B, npoint, C]
        # new_points: sampled points data, [B, npoint, nsample, C+D]
        # print(new_points.shape)

        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points =  F.relu(bn(conv(new_points)))

        new_points = torch.max(new_points, 2)[0] # max at dimension of nsample 
        new_xyz = new_xyz.permute(0, 2, 1)
        return new_xyz, new_points


class PointNetSetAbstractionMsg(nn.Module):
    def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
        super(PointNetSetAbstractionMsg, self).__init__()
        self.npoint = npoint
        self.radius_list = radius_list
        self.nsample_list = nsample_list
        self.conv_blocks = nn.ModuleList()
        self.bn_blocks = nn.ModuleList()
        for i in range(len(mlp_list)):
            convs = nn.ModuleList()
            bns = nn.ModuleList()
            last_channel = in_channel + 3
            for out_channel in mlp_list[i]:
                convs.append(nn.Conv2d(last_channel, out_channel, 1))
                bns.append(nn.BatchNorm2d(out_channel))
                last_channel = out_channel
            self.conv_blocks.append(convs)
            self.bn_blocks.append(bns)

    def forward(self, xyz, points):
        """
        Input:
            xyz: input points position data, [B, C, N]
            points: input points data, [B, D, N]
        Return:
            new_xyz: sampled points position data, [B, C, S]
            new_points_concat: sample points feature data, [B, D', S]
        """
        xyz = xyz.permute(0, 2, 1)
        if points is not None:
            points = points.permute(0, 2, 1)

        B, N, C = xyz.shape
        S = self.npoint
        new_xyz = index_points(xyz, farthest_point_sample(xyz, S))
        new_points_list = []
        for i, radius in enumerate(self.radius_list):
            K = self.nsample_list[i]
            group_idx = query_ball_point(radius, K, xyz, new_xyz)
            grouped_xyz = index_points(xyz, group_idx)
            grouped_xyz -= new_xyz.view(B, S, 1, C)
            if points is not None:
                grouped_points = index_points(points, group_idx)
                grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
            else:
                grouped_points = grouped_xyz

            grouped_points = grouped_points.permute(0, 3, 2, 1)  # [B, D, K, S]
            for j in range(len(self.conv_blocks[i])):
                conv = self.conv_blocks[i][j]
                bn = self.bn_blocks[i][j]
                grouped_points =  F.relu(bn(conv(grouped_points)))
            new_points = torch.max(grouped_points, 2)[0]  # [B, D', S]
            new_points_list.append(new_points)

        new_xyz = new_xyz.permute(0, 2, 1)
        new_points_concat = torch.cat(new_points_list, dim=1)
        return new_xyz, new_points_concat


class PointNetFeaturePropagation(nn.Module):
    def __init__(self, in_channel, mlp):
        super(PointNetFeaturePropagation, self).__init__()
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm1d(out_channel))
            last_channel = out_channel

    def forward(self, xyz1, xyz2, points1, points2):
        """
        Input:
            xyz1: input points position data, [B, C, N]
            xyz2: sampled input points position data, [B, C, S]
            points1: input points data, [B, D, N]
            points2: input points data, [B, D, S]
        Return:
            new_points: upsampled points data, [B, D', N]
        """
        xyz1 = xyz1.permute(0, 2, 1)
        xyz2 = xyz2.permute(0, 2, 1)

        points2 = points2.permute(0, 2, 1)
        B, N, C = xyz1.shape
        _, S, _ = xyz2.shape

        if S == 1:
            interpolated_points = points2.repeat(1, N, 1)
        else:
            dists = square_distance(xyz1, xyz2)
            dists, idx = dists.sort(dim=-1)
            dists, idx = dists[:, :, :3], idx[:, :, :3]  # [B, N, 3]

            dist_recip = 1.0 / (dists + 1e-8)
            norm = torch.sum(dist_recip, dim=2, keepdim=True)
            weight = dist_recip / norm
            interpolated_points = torch.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2)

        if points1 is not None:
            points1 = points1.permute(0, 2, 1)
            new_points = torch.cat([points1, interpolated_points], dim=-1)
        else:
            new_points = interpolated_points

        new_points = new_points.permute(0, 2, 1)
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points = F.relu(bn(conv(new_points)))
        return new_points

pointnet2_cls.py

import torch.nn as nn
import torch
import torch.nn.functional as F
from .pointnet2_utils import PointNetSetAbstraction
import logging

class get_model(nn.Module):
    def __init__(self,num_class,normal_channel=True,pretrained=False,load_path=None):
        super(get_model, self).__init__()
        in_channel = 6 if normal_channel else 3
        self.normal_channel = normal_channel
        self.sa1 = PointNetSetAbstraction(npoint=512, radius=0.2, nsample=64, in_channel=in_channel, mlp=[64, 64, 128], group_all=False)
        self.sa2 = PointNetSetAbstraction(npoint=256, radius=0.4, nsample=64, in_channel=128+3, mlp=[128, 128, 256], group_all=False)
        self.sa3 = PointNetSetAbstraction(npoint=None, radius=None, nsample=None, in_channel=256+3, mlp=[256, 512, 1024], group_all=True)
        self.fc1 = nn.Linear(1024, 512)
        self.bn1 = nn.BatchNorm1d(512)
        self.drop1 = nn.Dropout(0.4)
        self.fc2 = nn.Linear(512, 256)
        self.bn2 = nn.BatchNorm1d(256)
        self.drop2 = nn.Dropout(0.4)
        self.fc3 = nn.Linear(256, num_class)

        if pretrained:
            if load_path is not None:
                self.load_state_dict(torch.load(load_path)['model_state_dict'])
                logging.info("=> done loading ")
            else:
                logging.info("=> no provided checkpoint path ")
            
    def forward(self, xyz):
        
        B, _, _ = xyz.shape
        xyz = xyz.transpose(2, 1)
        print(xyz.shape)
        if self.normal_channel:
            norm = xyz[:, 3:, :]
            xyz = xyz[:, :3, :]
        else:
            norm = None
            
        l1_xyz, l1_points = self.sa1(xyz, norm)
        # print(l1_xyz.shape, l1_points.shape)
        l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
        # print(l2_xyz.shape, l2_points.shape)
        l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
        # print(l3_xyz.shape, l3_points.shape)
        x = l3_points.view(B, 1024)
        # print(x)
        x = self.drop1(F.relu(self.bn1(self.fc1(x))))
        x = self.drop2(F.relu(self.bn2(self.fc2(x))))
        x = self.fc3(x)
        # x = F.log_softmax(x, -1)


        return x, l3_points

def pointnetv2_cls(num_class,normal_channel=False,pretrained=False,load_path=None):
    return get_model(num_class,normal_channel=normal_channel,pretrained=pretrained,load_path=load_path)

class get_cls_loss(nn.Module):
    def __init__(self):
        super(get_cls_loss, self).__init__()

    def forward(self, pred, target, trans_feat):
        total_loss = F.nll_loss(pred, target)

        return total_loss

pointnet_v2_classification_int8.json

{
    "model": "pointnetv2",
    "pretrained": false,
    "dataset_dir":"/data1/modelnet/modelnet40_normal_resampled",
    "input_info": {
      "sample_size": [1, 1024, 3]
    },
    "num_classes": 40,
    "batch_size" : 64,
    "epochs": 80,
    "optimizer": {
        "type": "Adam",
        "base_lr": 0.00001,
        "schedule_type": "multistep",
        "steps": [
            5
        ]
    },
    "compression": {
        "algorithm": "quantization",
        "initializer": {
            "range": {
                "num_init_samples": 2560
            }
        }
    }
}

[AlexanderDokuchaev]

@zbnlala unfortunately this issue is not reproduced.

Could you check that issue is not reproduced without nncf?