p09823145
/
InfiniTensor
forked from jiuyuan/InfiniTensor
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

fix: fix re-dataMalloc for weight tensor and use of naive allocator

This commit is contained in:
kilinchange 2023-12-29 17:27:36 +08:00
parent 935b465cf2
commit dc6befb549
3 changed files with 19 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
include/core/tensor.h
View File

@ -42,8 +42,16 @@ class TensorObj : public TensorBaseObj {
bool isOutput() const { return tensorType == TensorType::output; } bool isOutput() const { return tensorType == TensorType::output; }
bool isOthers() const { return tensorType == TensorType::others; } bool isOthers() const { return tensorType == TensorType::others; }
void setWeight() { tensorType = TensorType::weight; } void setWeight() { tensorType = TensorType::weight; }
void setInput() { tensorType = TensorType::input; } void setInput() {
void setOutput() { tensorType = TensorType::output; } if (!this->isWeight()) {
tensorType = TensorType::input;
}
}
void setOutput() {
if (!this->isWeight()) {
tensorType = TensorType::output;
}
}
string tensorTypeToString() const { string tensorTypeToString() const {
switch (tensorType) { switch (tensorType) {
case TensorType::weight: case TensorType::weight:

9
pyinfinitensor/src/pyinfinitensor/onnx.py
View File

@ -37,7 +37,7 @@ class OnnxStub:
It can be generated from an Onnx model object. It can be generated from an Onnx model object.
""" """
def __init__(self, model: ModelProto, runtime): def __init__(self, model: ModelProto, runtime, use_naive_allocator: bool = False):
# We use some user-defined operators for distributed inference # We use some user-defined operators for distributed inference
try: try:
# onnx simplifier performs inplace simplify # onnx simplifier performs inplace simplify
@ -54,6 +54,7 @@ class OnnxStub:
self.tensors: Dict[str, backend.Tensor] = {} self.tensors: Dict[str, backend.Tensor] = {}
self.tensor_node_map: Dict[str, str] = {} self.tensor_node_map: Dict[str, str] = {}
self.initializer: Dict[int, TensorProto] = {} self.initializer: Dict[int, TensorProto] = {}
self.use_naive_allocator: bool = use_naive_allocator
try: try:
model = infer_shapes(model) model = infer_shapes(model)
except: except:
@ -885,7 +886,7 @@ class OnnxStub:
################################ ################################
# Allocate memory space for data # Allocate memory space for data
################################ ################################
self.handler.data_malloc() self.handler.data_malloc(self.use_naive_allocator)
################################# #################################
# Copy in data to tensor objects # Copy in data to tensor objects
@ -1223,7 +1224,7 @@ class OnnxStub:
return ctx.build(name) return ctx.build(name)
def init(self) -> None: def init(self) -> None:
self.handler.data_malloc() self.handler.data_malloc(self.use_naive_allocator)
def optimize(self) -> None: def optimize(self) -> None:
self.handler.optimize() self.handler.optimize()
@ -1239,7 +1240,7 @@ class OnnxStub:
oldTensor = self.inputs[oldInput] oldTensor = self.inputs[oldInput]
self.handler.change_shape(newInput, oldTensor.fuid()) self.handler.change_shape(newInput, oldTensor.fuid())
self.handler.shape_infer() self.handler.shape_infer()
self.handler.data_malloc() self.handler.data_malloc(self.use_naive_allocator)
def getShape(self, name: str) -> List[int]: def getShape(self, name: str) -> List[int]:
if name in self.inputs: if name in self.inputs:

5
src/core/graph.cc
View File

@ -167,7 +167,10 @@ void GraphObj::dataMalloc(bool useNaiveAllocator, size_t memPoolSize) {
// note: behavior may not match running in non-naive mode, and it may // note: behavior may not match running in non-naive mode, and it may
// not reproduce the bug // not reproduce the bug
for (auto &tensor : tensors) { for (auto &tensor : tensors) {
tensor->dataMalloc(); if (!tensor->isWeight() ||
(tensor->isWeight() && !weightAllocated)) {
tensor->dataMalloc();
}
} }
return; return;
} }