TensorList.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#pragma once

#include <cstddef>
#include <memory>
#include <string>

#include "Open3D/Core/Blob.h"
#include "Open3D/Core/Device.h"
#include "Open3D/Core/Dtype.h"
#include "Open3D/Core/ShapeUtil.h"
#include "Open3D/Core/SizeVector.h"
#include "Open3D/Core/Tensor.h"
#include "Open3D/Core/TensorKey.h"
namespace open3d {

class TensorList {
public:
    TensorList(const SizeVector& shape,
               Dtype dtype,
               const Device& device = Device("CPU:0"),
               const int64_t& size = 0);

    TensorList(const std::vector<Tensor>& tensors,
               const Device& device = Device("CPU:0"));

    TensorList(const std::initializer_list<Tensor>& tensors,
               const Device& device = Device("CPU:0"));

    template <class InputIterator>
    TensorList(InputIterator first,
               InputIterator last,
               const Device& device = Device("CPU:0"))
        : device_(device) {
        ConstructFromIterators(first, last);
    }

    TensorList(const Tensor& internal_tensor, bool inplace = true);

    static TensorList FromTensor(const Tensor& tensor, bool inplace = false);

    TensorList(const TensorList& other);

    void CopyFrom(const TensorList& other);

    TensorList& operator=(const TensorList& other) &;

    void ShallowCopyFrom(const TensorList& other);

    Tensor AsTensor() const;

    void Resize(int64_t n);

    void PushBack(const Tensor& tensor);

    static TensorList Concatenate(const TensorList& a, const TensorList& b);

    TensorList operator+(const TensorList& other) const {
        return Concatenate(*this, other);
    }

    void Extend(const TensorList& other);

    TensorList& operator+=(const TensorList& other) {
        Extend(other);
        return *this;
    }

    Tensor operator[](int64_t index) const;

    void Clear();

    std::string ToString() const;

    SizeVector GetShape() const { return shape_; }

    Device GetDevice() const { return device_; }

    Dtype GetDtype() const { return dtype_; }

    int64_t GetSize() const { return size_; }

    int64_t GetReservedSize() const { return reserved_size_; }

    const Tensor& GetInternalTensor() const { return internal_tensor_; }

protected:
    // The shared internal constructor for iterators.
    template <class InputIterator>
    void ConstructFromIterators(InputIterator first, InputIterator last) {
        int64_t size = std::distance(first, last);
        if (size == 0) {
            utility::LogError(
                    "Empty input tensors cannot initialize a TensorList.");
        }

        // Infer size and reserved_size
        size_ = size;
        reserved_size_ = ReserveSize(size_);

        // Infer shape
        shape_ = std::accumulate(
                std::next(first), last, first->GetShape(),
                [](const SizeVector shape, const Tensor& tensor) {
                    return shape_util::BroadcastedShape(std::move(shape),
                                                        tensor.GetShape());
                });

        // Infer dtype
        dtype_ = first->GetDtype();
        bool dtype_consistent = std::accumulate(
                std::next(first), last, true,
                [&](bool same_type, const Tensor& tensor) {
                    return same_type && (dtype_ == tensor.GetDtype());
                });
        if (!dtype_consistent) {
            utility::LogError(
                    "Inconsistent tensor dtypes in tensors are not supported "
                    "in TensorList.");
        }

        // Construct internal tensor
        SizeVector expanded_shape = ExpandFrontDim(shape_, reserved_size_);
        internal_tensor_ = Tensor(expanded_shape, dtype_, device_);

        // Assign tensors
        size_t i = 0;
        for (auto iter = first; iter != last; ++iter, ++i) {
            internal_tensor_[i] = *iter;
        }
    }

    void ExpandTensor(int64_t new_reserved_size);

    static SizeVector ExpandFrontDim(const SizeVector& shape,
                                     int64_t new_dim_size = 1);

    int64_t ReserveSize(int64_t n);

protected:
    SizeVector shape_;

    Dtype dtype_;
    Device device_;

    int64_t reserved_size_ = 0;

    int64_t size_ = 0;

    Tensor internal_tensor_;
};
}  // namespace open3d