search

NVIDIA-AI-IOT / tf_trt_models

TensorFlow models accelerated with NVIDIA TensorRT
BSD 3-Clause "New" or "Revised" License
682 stars 245 forks source link

the function `build_classification_graph` can use my own defined model ? #12

Open PapaMadeleine2022 opened 5 years ago

PapaMadeleine2022 commented 5 years ago

hello, I have a question, the function build_classification_graph can use my own model by replace the parameter model to my own network path? in which the net is defined as follows:

class MyLeNet(Network):
    def setup(self):
        (self.feed('data')
             .conv(5, 5, 20, 1, 1, padding='VALID', relu=False, name='conv1')
             .max_pool(2, 2, 2, 2, name='pool1')
             .conv(5, 5, 50, 1, 1, padding='VALID', relu=False, name='conv2')
             .max_pool(2, 2, 2, 2, name='pool2')
             .fc(500, name='ip1')
             .fc(10, relu=False, name='ip2')
             .softmax(name='prob'))