import sklearn
import torch
import numpy as np
from mindspore import context
from src.compare import compare
from module.encoder_pt import cnnEncoder as pt_encoder
from module.encoder_ms import cnnEncoder as ms_encoder
# from module.ms_mse_loss import MyMseLoss as ms_encoder
# from module.pt_mse_loss import MyMseLoss as pt_encoder

context.set_context(device_target="Ascend", device_id=7)

if __name__ == "__main__":
    # denfine your own net
    PtNet = pt_encoder()
    MsNet = ms_encoder()
    
    cp = compare(ptmodule=PtNet, 
                 msmodule=MsNet, 
                 module_type="net",  # "net" or "loss"
                 input_shape=[[1,1,28,28]],  # if 2 inputs, format is [[],[]]
                 init_mode="random",  # "random" or "ones"
                 input_num=1,  # 1 or 2, modify input_shape at the same time
                 print_result=True
                 )
    cp.start_compare()