import numpy as np
import mindspore.nn as msnn
import mindspore.ops as ops

class hetero_loss(msnn.Cell):
    def __init__(self, margin=0.1, dist_type='l2'):
        super(hetero_loss, self).__init__()
        self.margin = margin
        self.dist_type = dist_type
        if dist_type == 'l2':
            self.dist = msnn.MSELoss(reduction='sum')
        if dist_type == 'cos':
            self.dist = msnn.CosineSimilarity(dim=0)
        if dist_type == 'l1':
            self.dist = msnn.L1Loss()

    def construct(self, feat1, feat2, label1, label2):
        feat_size = feat1.shape[1]
        feat_num = feat1.shape[0]
        label_num = len(ops.unique(label1))
        feat1 = feat1.chunk(label_num, 0)
        feat2 = feat2.chunk(label_num, 0)
        # loss = Variable(.cuda())
        dist = 0
        for i in range(label_num):
            # f2 = tensor(feat2[i])
            # f1 = tensor(feat1[i])
            center1 = ops.mean(feat1[i],  axis=0)
            center2 = ops.mean(feat2[i], axis=0)
            if self.dist_type == 'l2' or self.dist_type == 'l1':
                if i == 0:

                    dist = 0 if 0 > self.dist(
                        center1, center2) - self.margin else self.dist(center1, center2)
                    # dist = max(0, self.dist(center1, center2) - self.margin)
                else:
                    num = 0 if 0 > self.dist(
                        center1, center2) - self.margin else self.dist(center1, center2)
                    dist += num
            elif self.dist_type == 'cos':
                if i == 0:
                    dist = 0 if 0 > 1-self.dist(center1, center2) - self.margin else 1-self.dist(center1, center2) - self.margin
                    # dist = max(0, 1-self.dist(center1, center2) - self.margin)
                else:
                    num =0 if 0 >  1-self.dist(center1, center2) - self.margin else 1-self.dist(center1, center2) - self.margin
                    dist += num

        return dist

==> Loading data..
Dataset sysu statistics:
  ------------------------------
  subset   | # ids | # images
  ------------------------------
  visible  |   395 |    22258
  thermal  |   395 |    11909
  ------------------------------
  query    |    96 |     3803
  gallery  |    96 |      301
  ------------------------------
==> Start Training...
  0%|             (48, 3, 288, 144) (48, 3, 288, 144) (48,) (48,)                                   | 0/60 [00:00<?, ?it/s]
training....?it/s]
&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
------start-----------
net output...........
labels000000000000
*******************
==========================
[ERROR] DEVICE(1130203,ffff85d4b010,python):2024-06-28-16:46:49.413.553 [mindspore/ccsrc/plugin/device/ascend/hal/hardware/ge_kernel_executor.cc:139] PrintQueryAclTypeErr] The current [Default/AvgPool3DGrad-op507] operator did not match any operator prototype library. The reason is:
The supported input and output data types for the current operator are:
InputDesc [0] support {int32,}
InputDesc [1] support {float16,float,double,}
OutputDesc [0] support {float16,float,double,}

But current operator's input and output data types is:
InputDesc [0] is Int32
InputDesc [1] is Float32
InputDesc [0] is Float32

[ERROR] KERNEL(1130203,fffa517fa0e0,python):2024-06-28-16:47:04.115.426 [mindspore/ccsrc/plugin/device/ascend/kernel/acl/acl_kernel_mod.cc:253] Launch] Kernel launch failed, msg: Acl compile and execute failed, op_type_:Pow

----------------------------------------------------
- Ascend Error Message:
----------------------------------------------------
E29999: Inner Error!
E29999  [GraphOpt][Trans][InsertTransByConcec] We do not support trans from NC1HWC0 to ND between Pow25_in_0 and Pow25.[FUNC:InsertTransOpByConcecutiveStrategy][FILE:trans_node_insertion.cc][LINE:545]
        TraceBack (most recent call last):
        [GraphOpt][Trans][Insert] Failed to insert format and dtype transfer op for graph online.[FUNC:InsertTransNodesForAllGraph][FILE:fe_graph_optimizer.cc][LINE:402]
        Call OptimizeOriginalGraphJudgeInsert failed, ret:-1, engine_name:AIcoreEngine, graph_name:online[FUNC:OptimizeOriginalGraphJudgeInsert][FILE:graph_optimize.cc][LINE:251]
        build graph failed, graph id:24, ret:-1[FUNC:BuildModelWithGraphId][FILE:ge_generator.cc][LINE:1615]
        [Build][SingleOpModel]call ge interface generator.BuildSingleOpModel failed. ge result = 4294967295[FUNC:ReportCallError][FILE:log_inner.cpp][LINE:161]
        [Build][Op]Fail to build op model[FUNC:ReportInnerError][FILE:log_inner.cpp][LINE:145]
        build op model failed, result = 500002[FUNC:ReportInnerError][FILE:log_inner.cpp][LINE:145]

(Please search "Ascend Error Message" at https://www.mindspore.cn for error code description)

----------------------------------------------------
- C++ Call Stack: (For framework developers)
----------------------------------------------------
mindspore/ccsrc/transform/acl_ir/acl_utils.cc:406 Run

[ERROR] DEVICE(1130203,fffa517fa0e0,python):2024-06-28-16:47:04.115.490 [mindspore/ccsrc/plugin/device/ascend/hal/hardware/ge_kernel_executor.cc:566] LaunchKernel] Launch kernel failed, kernel full name: Default/Pow-op527
0it [01:07, ?it/s]
  0%|                                                                                               | 0/60 [01:07<?, ?it/s]
Traceback (most recent call last):
  File "/home/ma-user/work/NNUN/DLRL_mindspore/train2.py", line 88, in <module>
    trainer1.train(epochs=60,trainset=trainset,color_pos=color_pos,thermal_pos=thermal_pos,per_img=per_img,\
  File "/home/ma-user/work/NNUN/DLRL_mindspore/trainer2.py", line 308, in train
    loss = self.train_single(input1_m, input2_m, label1_m, label2_m)
  File "/home/ma-user/work/NNUN/DLRL_mindspore/trainer2.py", line 139, in train_single
    loss, grads = self.value_and_grad(input1, input2, label1, label2)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/mindspore/ops/composite/base.py", line 625, in after_grad
    return grad_(fn_, weights)(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/mindspore/common/api.py", line 121, in wrapper
    results = fn(*arg, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/mindspore/ops/composite/base.py", line 602, in after_grad
    out = _pynative_executor()
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/mindspore/common/api.py", line 1147, in __call__
    return self._executor()
RuntimeError: Launch kernel failed, name:Default/Pow-op527

----------------------------------------------------
- Ascend Error Message:
----------------------------------------------------
E29999: Inner Error!
E29999  [GraphOpt][Trans][InsertTransByConcec] We do not support trans from NC1HWC0 to ND between Pow25_in_0 and Pow25.[FUNC:InsertTransOpByConcecutiveStrategy][FILE:trans_node_insertion.cc][LINE:545]
        TraceBack (most recent call last):
        [GraphOpt][Trans][Insert] Failed to insert format and dtype transfer op for graph online.[FUNC:InsertTransNodesForAllGraph][FILE:fe_graph_optimizer.cc][LINE:402]
        Call OptimizeOriginalGraphJudgeInsert failed, ret:-1, engine_name:AIcoreEngine, graph_name:online[FUNC:OptimizeOriginalGraphJudgeInsert][FILE:graph_optimize.cc][LINE:251]
        build graph failed, graph id:24, ret:-1[FUNC:BuildModelWithGraphId][FILE:ge_generator.cc][LINE:1615]
        [Build][SingleOpModel]call ge interface generator.BuildSingleOpModel failed. ge result = 4294967295[FUNC:ReportCallError][FILE:log_inner.cpp][LINE:161]
        [Build][Op]Fail to build op model[FUNC:ReportInnerError][FILE:log_inner.cpp][LINE:145]
        build op model failed, result = 500002[FUNC:ReportInnerError][FILE:log_inner.cpp][LINE:145]

(Please search "Ascend Error Message" at https://www.mindspore.cn for error code description)

----------------------------------------------------
- C++ Call Stack: (For framework developers)
----------------------------------------------------
mindspore/ccsrc/runtime/pynative/run_op_helper.cc:993 LaunchKernelsDynamic

Error in atexit._run_exitfuncs:
RuntimeError: Call aclrtDestroyStream failed, ret = 0x107003

----------------------------------------------------
- Ascend Error Message:
----------------------------------------------------
EE9999: Inner Error!
EE9999  Stream destroy failed, stream is not in current ctx.[FUNC:StreamDestroy][FILE:api_impl.cc][LINE:874]
        TraceBack (most recent call last):
        rtStreamDestroy execute failed, reason=[stream not in current context][FUNC:FuncErrorReason][FILE:error_message_manage.cc][LINE:50]
        destroy stream failed, runtime result = 107003[FUNC:ReportCallError][FILE:log_inner.cpp][LINE:161]

(Please search "Ascend Error Message" at https://www.mindspore.cn for error code description)
E29999: Inner Error!
E29999  [GraphOpt][Trans][InsertTransByConcec] We do not support trans from NC1HWC0 to ND between Pow25_in_0 and Pow25.[FUNC:InsertTransOpByConcecutiveStrategy][FILE:trans_node_insertion.cc][LINE:545]
        TraceBack (most recent call last):
        [GraphOpt][Trans][Insert] Failed to insert format and dtype transfer op for graph online.[FUNC:InsertTransNodesForAllGraph][FILE:fe_graph_optimizer.cc][LINE:402]
        Call OptimizeOriginalGraphJudgeInsert failed, ret:-1, engine_name:AIcoreEngine, graph_name:online[FUNC:OptimizeOriginalGraphJudgeInsert][FILE:graph_optimize.cc][LINE:251]
        build graph failed, graph id:24, ret:-1[FUNC:BuildModelWithGraphId][FILE:ge_generator.cc][LINE:1615]
        [Build][SingleOpModel]call ge interface generator.BuildSingleOpModel failed. ge result = 4294967295[FUNC:ReportCallError][FILE:log_inner.cpp][LINE:161]
        [Build][Op]Fail to build op model[FUNC:ReportInnerError][FILE:log_inner.cpp][LINE:145]
        build op model failed, result = 500002[FUNC:ReportInnerError][FILE:log_inner.cpp][LINE:145]

(Please search "Ascend Error Message" at https://www.mindspore.cn for error code description)

----------------------------------------------------
- C++ Call Stack: (For framework developers)
----------------------------------------------------
mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_kernel_load.cc:366 FreeDeviceMemory

Segmentation fault (core dumped)
(MindSpore) [ma-user DLRL_mindspore]$