Package: tensorneko
TensorNekoTensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.InstallpipinstalltensornekoTo use the library without PyTorch and PyTorch Lightning, you can install the util library (support Python 3.7 ~ 3.10 with limited features) with following command.pipinstalltensorneko_utilNeko Layers, Modules and ArchitecturesBuild an MLP with linear layers. The activation and normalization will be placed in the hidden layers.784 -> 1024 -> 512 -> 10importtensornekoasnekoimporttorch.nnmlp=neko.module.MLP(neurons=[784,1024,512,10],build_activation=torch.nn.ReLU,build_normalization=[lambda:torch.nn.BatchNorm1d(1024),lambda:torch.nn.BatchNorm1d(512)],dropout_rate=0.5)Build a Conv2d with activation and normalization.importtensornekoasnekoimporttorch.nnconv2d=neko.layer.Conv2d(in_channels=256,out_channels=1024,kernel_size=(3,3),padding=(1,1),build_activation=torch.nn.ReLU,build_normalization=lambda:torch.nn.BatchNorm2d(256),normalization_after_activation=False)All architectures, modules and layersLayers:AggregationConcatenateConv,Conv1d,Conv2d,Conv3dGaussianNoiseImageAttention,SeqAttentionMaskedConv2d,MaskedConv2dA,MaskedConv2dBLinearLogPatchEmbedding2dPositionalEmbeddingReshapeStackVectorQuantizerModules:DenseBlockInceptionModuleMLPResidualBlockandResidualModuleAttentionModule,TransformerEncoderBlockandTransformerEncoderGatedConvArchitectures:AutoEncoderGANWGANVQVAENeko modulesAlltensorneko.layerandtensorneko.moduleareNekoModule. They can be used infn.pypipe operation.fromtensorneko.layerimportLinearfromtorch.nnimportReLUimporttorchlinear0=Linear(16,128,build_activation=ReLU)linear1=Linear(128,1)f=linear0>>linear1print(f(torch.rand(16)).shape)# torch.Size([1])Neko IOEasily load and save different modal data.importtensornekoasnekofromtensorneko.ioimportjson_datafromtypingimportList# read video (Temporal, Channel, Height, Width)video_tensor,audio_tensor,video_info=neko.io.read.video("path/to/video.mp4")# write videoneko.io.write.video("path/to/video.mp4",video_tensor,video_info.video_fps,audio_tensor,video_info.audio_fps)# read audio (Channel, Temporal)audio_tensor,sample_rate=neko.io.read.audio("path/to/audio.wav")# write audioneko.io.write.audio("path/to/audio.wav",audio_tensor,sample_rate)# read image (Channel, Height, Width) with float value in range [0, 1]image_tensor=neko.io.read.image("path/to/image.png")# write imageneko.io.write.image("path/to/image.png",image_tensor)neko.io.write.image("path/to/image.jpg",image_tensor)# read plain texttext_string=neko.io.read.text("path/to/text.txt")# write plain textneko.io.write.text("path/to/text.txt",text_string)# read json as python dict or listjson_dict=neko.io.read.json("path/to/json.json")# read json as an object@json_dataclassJsonData:x:inty:intjson_obj:List[JsonData]=neko.io.read.json("path/to/json.json",cls=List[JsonData])# write json from python dict/list or json_data decorated objectneko.io.write.json("path/to/json.json",json_dict)neko.io.write.json("path/to/json.json",json_obj)Besides, the read/write formatandpicklefiles is also supported.Neko preprocessingimporttensornekoasneko# A video tensor with (120, 3, 720, 1280)video=neko.io.read.video("example/video.mp4").video# Get a resized tensor with (120, 3, 256, 256)resized_video=neko.preprocess.resize_video(video,(256,256))All preprocessing utilsresize_videoresize_imagepadding_videopadding_audiocrop_with_paddingframes2videoifffmpegis available, you can use below ffmpeg wrappers.video2framesmerge_video_audioresample_video_fpsmp32wavNeko VisualizationVariable Web WatcherStart a web server to watch the variable status when the program (e.g. training, inference, data preprocessing) is running.importtimefromtensorneko.visualization.watcherimport*data_list=...# a list of datadefpreprocessing(d):...# initialize the componentspb=ProgressBar("Processing",total=len(data_list))logger=Logger("Log message")var=Variable("Some Value",0)line_chart=LineChart("Line Chart",x_label="x",y_label="y")view=View("Data preprocessing").add_all()t0=time.time()# open server when the code block in running.withServer(view,port=8000):fori,datainenumerate(data_list):preprocessing(data)# do some processing herex=time.time()-t0# time since the start of the programy=i# processed number of dataline_chart.add(x,y)# add to the line chartlogger.log("Some messages")# log messages to the servervar.value=...# keep tracking a variablepb.add(1)# update the progress bar by add 1When the script is running, go to127.0.0.1:8000to keep tracking the status.Tensorboard ServerSimply run tensorboard server in Python script.importtensornekoasnekowithneko.visualization.tensorboard.Server(port=6006):trainer.fit(model,dm)Matplotlib wrappersDisplay an image of (C, H, W) shape byplt.imshowwrapper.importtensornekoasnekoimportmatplotlib.pyplotaspltimage_tensor=...# an image tensor with shape (C, H, W)neko.visualization.matplotlib.imshow(image_tensor)plt.show()Predefined colorsSeveral aesthetic colors are predefined.importtensornekoasnekoimportmatplotlib.pyplotasplt# use with matplotlibplt.plot(...,color=neko.visualization.Colors.RED)# the palette for seaborn is also availablefromtensorneko_util.visualization.seabornimportpaletteimportseabornassnssns.set_palette(palette)Neko ModelBuild and train a simple model for classifying MNIST with MLP.fromtypingimportOptional,Union,Sequence,Dict,Listimporttorch.nnfromtorchimportTensorfromtorch.optimimportAdamfromtorchmetricsimportAccuracyfromlightning.pytorch.callbacksimportModelCheckpointimporttensornekoasnekofromtensorneko.utilimportget_activation,get_lossclassMnistClassifier(neko.NekoModel):def__init__(self,name:str,mlp_neurons:List[int],activation:str,dropout_rate:float,loss:str,learning_rate:float,weight_decay:float):super().__init__(name)self.weight_decay=weight_decayself.learning_rate=learning_rateself.flatten=torch.nn.Flatten()self.mlp=neko.module.MLP(neurons=mlp_neurons,build_activation=get_activation(activation),dropout_rate=dropout_rate)self.loss_func=get_loss(loss)()self.acc_func=Accuracy()defforward(self,x):# (batch, 28, 28)x=self.flatten(x)# (batch, 768)x=self.mlp(x)# (batch, 10)returnxdeftraining_step(self,batch:Optional[Union[Tensor,Sequence[Tensor]]]=None,batch_idx:Optional[int]=None,optimizer_idx:Optional[int]=None,hiddens:Optional[Tensor]=None)->Dict[str,Tensor]:x,y=batchlogit=self(x)prob=logit.sigmoid()loss=self.loss_func(logit,y)acc=self.acc_func(prob.max(dim=1)[1],y)return{"loss":loss,"acc":acc}defvalidation_step(self,batch:Optional[Union[Tensor,Sequence[Tensor]]]=None,batch_idx:Optional[int]=None,dataloader_idx:Optional[int]=None)->Dict[str,Tensor]:x,y=batchlogit=self(x)prob=logit.sigmoid()loss=self.loss_func(logit,y)acc=self.acc_func(prob.max(dim=1)[1],y)return{"loss":loss,"acc":acc}defconfigure_optimizers(self):optimizer=Adam(self.parameters(),lr=self.learning_rate,betas=(0.5,0.9),weight_decay=self.weight_decay)return{"optimizer":optimizer}model=MnistClassifier("mnist_mlp_classifier",[784,1024,512,10],"ReLU",0.5,"CrossEntropyLoss",1e-4,1e-4)dm=...# The MNIST datamodule from PyTorch Lightningtrainer=neko.NekoTrainer(log_every_n_steps=100,gpus=1,logger=model.name,precision=32,callbacks=[ModelCheckpoint(dirpath="./ckpt",save_last=True,filename=model.name+"-{epoch}-{val_acc:.3f}",monitor="val_acc",mode="max")])trainer.fit(model,dm)Neko CallbacksSome simple but useful pytorch-lightning callbacks are provided.DisplayMetricsCallbackEarlyStoppingLR: Early stop training when learning rate reaches threshold.Neko Notebook HelpersHere are some helper functions to better interact with Jupyter Notebook.importtensornekoasneko# display a videoneko.notebook.display.video("path/to/video.mp4")# display an audioneko.notebook.display.audio("path/to/audio.wav")# display a code fileneko.notebook.display.code("path/to/code.java")Neko Debug ToolsGet the default values fromArgumentParserargs. It's convenient to use this in the notebook.fromargparseimportArgumentParserfromtensorneko.debugimportget_parser_default_argsparser=ArgumentParser()parser.add_argument("integers",type=int,nargs="+",default=[1,2,3])parser.add_argument("--sum",dest="accumulate",action="store_const",const=sum,default=max)args=get_parser_default_args(parser)print(args.integers)# [1, 2, 3]print(args.accumulate)# <function sum at ...>Neko EvaluationSome metrics function for evaluation are provided.iou_1diou_2dpsnr_videopsnr_imagessim_videossim_imageNeko UtilitiesMisc functions__: The arguments to pipe operator. (Inspired fromfn.py)fromtensorneko.utilimport__,_result=__(20)>>(_+1)>>(_*2)>>__.getprint(result)# 42SeqandStream: A collection wrapper for method chaining with concurrent supporting.fromtensorneko.utilimportSeq,Stream,_fromtensorneko_util.backend.parallelimportParallelType# using method chainingseq=Seq.of(1,2,3).map(_+1).filter(_%2==0).map(_*2).take(2).to_list()# return [4, 8]# using bit shift operator to chain the sequenceseq=Seq.of(1,2,3)<<Seq.of(2,3,4)<<[3,4,5]# return Seq(1, 2, 3, 2, 3, 4, 3, 4, 5)# run concurrent with `for_each` for Streamif__name__=='__main__':Stream.of(1,2,3,4).for_each(print,progress_bar=True,parallel_type=ParallelType.PROCESS)Option: A monad for dealing with data.fromtensorneko.utilimportreturn_option@return_optiondefget_data():ifsome_condition:return1else:returnNonedefprocess_data(n:int):ifcondition(n):returnnelse:returnNonedata=get_data()data=data.map(process_data).get_or_else(-1)# if the response is None, return -1Eval: A monad for lazy evaluation.fromtensorneko.utilimportEval@Eval.alwaysdefcall_by_name_var():return42@Eval.laterdefcall_by_need_var():return43@Eval.nowdefcall_by_value_var():return44print(call_by_name_var.value)# 42ReactiveThis library provides event bus based reactive tools. The API integrates the Python type annotation syntax.# useful decorators for default event busfromtensorneko.utilimportsubscribe# Event base typefromtensorneko.utilimportEventclassLogEvent(Event):def__init__(self,message:str):self.message=message# the event argument should be annotated correctly@subscribe# run in the main threaddeflog_information(event:LogEvent):print(event.message)@subscribe.thread# run in a new threaddeflog_information_thread(event:LogEvent):print(event.message,"in another thread")@subscribe.coro# run with asyncasyncdeflog_information_async(event:LogEvent):print(event.message,"async")@subscribe.process# run in a new processdeflog_information_process(event:LogEvent):print(event.message,"in a new process")if__name__=='__main__':# emit an event, and then the event handler will be invoked# The sequential order is not guaranteedLogEvent("Hello world!")# one possible output:# Hello world! in another thread# Hello world! async# Hello world!# Hello world! in a new processMultiple Dispatchdispatch: Multi-dispatch implementation for Python.To my knowledge, 3 popular multi-dispatch libraries still have critical limitations.plumdoesn't support static methods,mutipledispatchdoesn't support Python type annotation syntax andmultimethoddoesn't support default argument. TensorNeko can do it all.fromtensorneko.utilimportdispatchclassDispatchExample:@staticmethod@dispatchdefgo()->None:print("Go0")@staticmethod@dispatchdefgo(x:int)->None:print("Go1")@staticmethod@dispatchdefgo(x:float,y:float=1.0)->None:print("Go2")@dispatchdefcome(x:int)->str:return"Come1"@dispatch.of(str)defcome(x)->str:return"Come2"MiscellaneousStringGetter: Get PyTorch class from string.importtensornekoasnekoactivation=neko.util.get_activation("leakyRelu")()Seed: The universal seed fornumpy,torchand Pythonrandom.fromtensorneko.utilimportSeedfromtorch.utils.dataimportDataLoader# set seed to 42 for all numpy, torch and python randomSeed.set(42)# Apply seed to parallel workers of DataLoaderDataLoader(train_dataset,batch_size=batch_size,num_workers=num_workers,worker_init_fn=Seed.get_loader_worker_init(),generator=Seed.get_torch_generator())Timer: A timer for measuring the time.fromtensorneko.utilimportTimerimporttime# use as a context manager with single timewithTimer():time.sleep(1)# use as a context manager with multiple segmentswithTimer()ast:time.sleep(1)t.time("sleep A")time.sleep(1)t.time("sleep B")time.sleep(1)# use as a decorator@Timer()deff():time.sleep(1)print("f")Singleton: A decorator to make a class as a singleton. Inspired from Scala/Kotlin.fromtensorneko.utilimportSingleton@SingletonclassMyObject:def__init__(self):self.value=0defadd(self,value):self.value+=valuereturnself.valueprint(MyObject.value)# 0MyObject.add(1)print(MyObject.value)# 1Besides, many miscellaneous functions are also provided.Functions list (intensorneko_util):generate_inf_seqcomposelistdirwith_printedifelsedict_addas_listidentitylist_to_dictget_tensorneko_util_pathFunctions list (intensorneko):reduce_dict_bysummarize_dict_bywith_printed_shapeis_bad_numcount_parameters