Add more easy-use function interface

[Jittor]

Currently, visualize a computation graph needs to export pkl file and load in jittorvis, it is not straightforward, we can do it in a better way:

import jittorvis
import jittor as jt
from jittor.models import resnet18

model = resnet18()
img = jt.rand((1,3,100,100))

jittorvis.visualize_model(model, img, ip=xxxx, port=xxxx, ...)

[swordsbird]

in the updated version, we can use jittorvis in the following way:

import jittor as jt
from jittor.models import resnet18
from jittorvis import server

model = resnet18()
with jt.flag_scope(trace_py_var=2, trace_var_data=1):
  output = model(input)
  output.sync()
  trace_data = jt.dump_trace_data()
  jt.clear_trace_data()
server.run(trace_data, port=5005, host='0.0.0.0')

In this way you don't need to export pkl file and reload in jittorvis. Also, we will introduce the function "visualize_model" in jittorvis soon.

[swordsbird]

Solved