microsoft / nn-Meter

A DNN inference latency prediction toolkit for accurately modeling and predicting the latency on diverse edge devices.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

KeyError in graph_tool

jtseven opened this issue · comments

Within the get_node_type method in graph_tool.py there is an if-clause to check, if a certain key exists in graph.keys(), but if it fails, accessing the key is still tried and will fail with a KeyError:

def get_node_type(self, name):
    if name in self.graph.keys() and "attr" in self.graph[name].keys():  # <-- Check if name is in keys
        return self.graph[name]["attr"]["type"]
    else:
        logging.info(name, self.graph[name])  # <-- Causes key error
    return None

Instead, something like that could be done:

logging.info(name, self.graph.get(name, "Name not in graph"))

I'm not sure if it is expected behavior that the name does not exist within the keys. If it's not, a different, more precise error message should be printed and an exception should be raised accordingly.