大规模可扩展知识蒸馏框架 Pantheon

Teacher

pantheon.Teacher() source

: The class defined for the teacher model. Generate knowledge data and transfer them to the student model.

Args:

  • out_path (str|None) - The path to dump knowledge data for offline mode.
  • out_port (int|None) - The IP port number to send out knowledge for online mode, should be unique when launching multiple teachers in the same node.

Return: An object of class Teacher

pantheon.Teacher.start() source

: Start teacher service, sychronize with student and launch the thread
to monitor commands from student.

Args: None

Return: None

pantheon.Teacher.send(data) source

: Send one data object to student.

Args:

  • data (Python data): - The data to be sent, can be any type of Python data object.

Return: None

pantheon.Teacher.recv() source

: Recieve one data object from student.

Args: None

Return:

  • The received data, can be any type of Python data object.

pantheon.Teacher.dump(knowledge) source

: Dump one batch knowledge data into the output file, only used in the offline mode.

Args:

  • knowledge (dict): - The knowledge data to be dumped.

Return: None

pantheon.Teacher.start_knowledge_service(feed_list, schema, program, reader_config, exe, buf_size=10, times=1) source

: Start the knowledge service to generate and transfer knowledge data. In GPU mode, the devices to execute knowledge prediction will be determined by the
environment variable FLAGS_selected_gpus, or by CUDA_VISIBLE_DEVICES if it is not set, and by CPU_NUM (default 1) in CPU mode. Only supported in static graph.

Args:

  • feed_list (list): - A list of feed Variables or their names for the .. code-block:: guess

    input teacher Program.

  • schema (dict): - A dict to specify keys and fetched Variables .. code-block:: guess

    to generate knowledge.

  • program (fluid.Program): - Inference Program of the teacher model.

  • reader_config (dict): - The config for data reader. Support all the three types of generators used by fluid.io.PyReader and fluid.io.DataLoader, and their configs contain the key-value pair of the generator type and a generator object, plus other necessary argument pairs. See the following:

    1. sample generator:
       reader_config={"sample_generator": some_sample_generator,
                      "batch_size": batch_size, "drop_last": drop_last}
       # drop_last set to True by default
    
    2) **sample list generator:**
    
       reader_config={"sample_list_generator": some_sample_list_generator}
    
    3) **batch generator:**
    
       reader_config={"batch_generator": some_batch_genrator}
    
    The trial to parse config will be in the order of 1) -> 3), and any other unrelated keys in these configs will be ignored.
    
  • exe (fluid.Executor): The executor to run the input program.

  • buf_size (int): The size of buffers for data reader and knowledge .. code-block:: guess

    writer on each device.

  • times (int): The maximum repeated serving times, default 1. Whenever .. code-block:: guess

    the public method get_knowledge_generator() in Student object called once, the serving times will be added one, until reaching the maximum and ending the service. Only valid in online mode, and will be ignored in offline mode.

Return: None

Examples:

import paddle
import paddle.fluid as fluid
from paddleslim.pantheon import Teacher

startup = fluid.Program()
program = fluid.Program()
with fluid.program_guard(program, startup):
    images = fluid.data(
            name='pixel', shape=[None, 3 * 32 * 32], dtype='float32')
    labels = fluid.data(name='label', shape=[None, 1], dtype='int64')
    logits = fluid.layers.fc(input=images, size=10)
    loss = fluid.layers.softmax_with_cross_entropy(logits, labels)

place = fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(startup)

train_reader = paddle.batch(
        paddle.dataset.cifar.train10(), batch_size=32)

teacher = Teacher(out_path="example_knowledge.dat", # offline mode
                  #out_port=5000                    # online mode
                  )
teacher.start()

teacher.start_knowledge_service(
    feed_list=[images, labels],
    schema={"logits": logits,
            "labels": labels},
    program=program,
    reader_config={"sample_list_generator": train_reader},
    exe=exe)
!!! note “Note”
This example should be run with the example of class Student.

Student

pantheon.Student(merge_strategy=None) source

: The class defined for the student model. Receive knowledge data from
teacher model and carry out knowledge merging.

Args:

  • merge_strategy (dict|None): - A dict whose keys are the common schemas shared by different teachers, and each corresponding value specifies the merging strategy for different schemas respectively, supporting sum and mean now.

Return: An object of class Student.

pantheon.Student.register_teacher(in_path=None, in_address=None) source

: Register one teacher model and assign the order number to it as its id, with the file path (offline mode) or IP address (online mode) that the teacher model writes knowledge data to.

Args:

  • in_path (str|None): The input file path. Default None.
  • in_address (str|None): The input IP address, in the format “<IP_address>:<IP_port>” (e.g. “127.0.0.1:8080”). Default None.

Return: None

pantheon.Student.start() source

: End teachers’ registration and synchronize with all of them.

Args: None

Return: None

pantheon.Student.send(self, data, teacher_ids=None) source

: Send data to teachers.

Args:

  • data (Python data): - A Python data object to be sent.
  • teacher_ids (list|None): - A list of teacher ids to send data. If set to None, send the data to all teachers. Default None.

Return: None

pantheon.Student.recv(teacher_id) source

: Receive data from one teacher.

Args:
  • teacher_id (int): - The id of teacher that receives data from.

Return:

  • The received data object.

pantheon.Student.get_knowledge_desc() source

: Get description for knowledge, including shape, data type and lod level for each schema.

Args: None

Return:

  • Knowledge description, which is a dict.

pantheon.Student.get_knowledge_qsize() source

: Get the real-time size of knowledge queue. If this size is denoted as
qsize, it means that there are qsize batch knowledge data already pushed into knowledge queue and waiting for the knowledge generator to pop out. It’s dynamic and limited up to 100, the capacity of the knowledge queue.

Args: None

Return:

  • The real-time size of knowledge queue.

pantheon.Student.get_knowledge_generator(batch_size, drop_last=False) source

: Get the generator for knowledge data, return None if last generator doesn’t finish yet.

Args:

  • batch_size (int): - The batch size of returned knowledge data.
  • drop_last (bool): - Whether to drop the last batch if its size is less than batch size.

Return:

  • The wrapper of knowledge data generator.

Examples:

from paddleslim.pantheon import Student

student = Student()

student.register_teacher(in_path="example_knowledge.dat",  # offline mode
                         #in_address="127.0.0.1:5000"      # online mode
                         )
student.start()

knowledge_desc = student.get_knowledge_desc()
data_generator = student.get_knowledge_generator(
    batch_size=128, drop_last=True)

# get knowledge data
for knowledge in data_generator():
    print("knowledge queue size: {}".format(student.get_knowledge_qsize()))

    # do something else
!!! note “Note”
This example should be run with the example of class Teacher.