Refactor hf/raven.

models.py

@@ -11,3 +11,4 @@ indexes = nload("/home/jkwiatkowski/all/dataset/arr/val_target.npy")
 
 folders = DataSetFromFolder("/home/jkwiatkowski/all/dataset/arr/RAVEN-10000-release/RAVEN-10000", file_type="dir")
 properties = DataSetFromFolder(folders[:], file_type="xml", extension="val")
+

raven_utils/depricated/old_raven.py

@@ -1,490 +0,0 @@
-from functools import partial
-
-import numpy as np
-from data_utils import take, EXIST, COR
-from data_utils.image import draw_images, add_text
-from data_utils.op import np_split
-from ml_utils import lu, dict_from_list2, filter_keys, none
-from data_utils import ops as K
-
-from config.constant import PROPERTY, TARGET, INPUTS
-# from raven_utils.render.rendering import render_panels
-
-RENDER_POSITIONS = [
-    [(0.5, 0.5, 1, 1)],
-    # ...
-    [(0.25, 0.25, 0.5, 0.5),
-     (0.25, 0.75, 0.5, 0.5),
-     (0.75, 0.25, 0.5, 0.5),
-     (0.75, 0.75, 0.5, 0.5)],
-    # ...
-    [(0.16, 0.16, 0.33, 0.33),
-     (0.16, 0.5, 0.33, 0.33),
-     (0.16, 0.83, 0.33, 0.33),
-     (0.5, 0.16, 0.33, 0.33),
-     (0.5, 0.5, 0.33, 0.33),
-     (0.5, 0.83, 0.33, 0.33),
-     (0.83, 0.16, 0.33, 0.33),
-     (0.83, 0.5, 0.33, 0.33),
-     (0.83, 0.83, 0.33, 0.33)],
-    # ...
-    [(0.5, 0.25, 0.5, 0.5)],
-    [(0.5, 0.75, 0.5, 0.5)],
-    # ...
-    [(0.25, 0.5, 0.5, 0.5)],
-    [(0.75, 0.5, 0.5, 0.5)],
-    # ...
-    [(0.5, 0.5, 1, 1)],
-    [(0.5, 0.5, 0.33, 0.33)],
-    # ...
-    [(0.5, 0.5, 1, 1)],
-    [(0.42, 0.42, 0.15, 0.15),
-     (0.42, 0.58, 0.15, 0.15),
-     (0.58, 0.42, 0.15, 0.15),
-     (0.58, 0.58, 0.15, 0.15)],
-    # ...
-
-]
-
-HORIZONTAL = "horizontal"
-VERTICAL = "vertical"
-
-NAMES = ['center_single',
-         'distribute_four',
-         'distribute_nine',
-         'in_center_single_out_center_single',
-         'in_distribute_four_out_center_single',
-         'left_center_single_right_center_single',
-         'up_center_single_down_center_single']
-
-PROPERTIES_NAMES = [
-    'Color',
-    'Size',
-    'Type',
-
-]
-PROPERTIES = dict_from_list2(PROPERTIES_NAMES, [10, 6, 5])
-ANGLE_MAX = 7
-
-PROPERTIES_NO = len(PROPERTIES)
-
-RULES_COMBINE = "Number/Position"
-
-RULES_ATTRIBUTES = [
-    "Number",
-    "Position",
-    "Color",
-    "Size",
-    "Type"
-]
-RULES_ATTRIBUTES_LEN = len(RULES_ATTRIBUTES)
-
-RULES_ATTRIBUTES_INDEX = dict_from_list2(RULES_ATTRIBUTES)
-
-RULES_TYPES = [
-    "Constant",
-    "Arithmetic",
-    "Progression",
-    "Distribute_Three"
-]
-RULES_TYPES_INDEX = dict_from_list2(RULES_TYPES)
-RULES_TYPES_LEN = len(RULES_ATTRIBUTES)
-
-GROUPS_NO = len(NAMES)
-ENTITY_NO = dict(zip(NAMES, [1, 4, 9, 2, 5, 2, 2]))
-ENTITY_SUM = sum(list(ENTITY_NO.values()))
-ENTITY_INDEX = np.concatenate([[0], np.cumsum(list(ENTITY_NO.values()))])
-ENTITY_INDEX_TARGET = ENTITY_INDEX + 1
-ENTITY_DICT = dict(zip(NAMES, ENTITY_INDEX_TARGET[:-1]))
-NAMES_ORDER = dict(zip(NAMES, np.arange(len(NAMES))))
-PROPERTIES_INDEXES = np.cumsum(np.array(list(ENTITY_NO.values())) * len(PROPERTIES))
-INDEX = np.concatenate([[0], PROPERTIES_INDEXES]) + ENTITY_SUM + 1  # +2 type and uniformity
-
-SECOND_LAYOUT = [i - 1 for i in [
-    ENTITY_DICT["in_center_single_out_center_single"] + 1,
-    ENTITY_DICT["in_distribute_four_out_center_single"] + 1,
-    ENTITY_DICT["in_distribute_four_out_center_single"] + 2,
-    ENTITY_DICT["in_distribute_four_out_center_single"] + 3,
-    ENTITY_DICT["left_center_single_right_center_single"] + 1,
-    ENTITY_DICT["up_center_single_down_center_single"] + 1
-]]
-
-FIRST_LAYOUT = list(set(range(ENTITY_SUM)) - set(SECOND_LAYOUT))
-LAYOUT_NO = 2
-
-START_INDEX = dict(zip(NAMES, INDEX[:-1]))
-END_INDEX = INDEX[-1]
-
-RULES_ATTRIBUTES_ALL_LEN = RULES_ATTRIBUTES_LEN * LAYOUT_NO
-UNIFORMITY_NO = 2
-UNIFORMITY_INDEX = END_INDEX + RULES_ATTRIBUTES_ALL_LEN
-
-FEATURE_NO = UNIFORMITY_INDEX + UNIFORMITY_NO
-MAPPING = {
-    "distribute_nine":
-        {0.16: 0,
-         0.5: 1,
-         0.83: 2},
-    "distribute_four":
-        {0.25: 0,
-         0.75: 1},
-    'in_distribute_four_out_center_single':
-        {0.42: 0,
-         0.58: 1}
-}
-MUL = {
-    "distribute_nine": 3,
-    "distribute_four": 2,
-    'in_distribute_four_out_center_single': 2
-}
-
-# SIZES = np.linspace(0.4, 0.9, 6)
-TYPES = ["triangle", "square", "pentagon", "hexagon", "circle"]
-# TYPES = ["triangle", "square", "pentagon", "circle", "circle"]
-SIZES = ["vs", "s", "m", "h", "vh", "e"]
-COLORS = ["vs", "s", "m", "h", "vh", "e"]
-# TYPES = ["", "", "circle", "hexagon", "square"]
-
-ENTITY_PROPERTIES_VALUES = list(PROPERTIES.values())
-ENTITY_PROPERTIES_KEYS = list(PROPERTIES.keys())
-ENTITY_PROPERTIES_NO = len(PROPERTIES)
-INDEX = dict(zip(PROPERTIES, np.array(ENTITY_PROPERTIES_VALUES) * ENTITY_SUM))
-ENTITY_PROPERTIES_SUM = sum(list(PROPERTIES.values()))
-
-OUTPUT_SIZE = ENTITY_SUM * ENTITY_PROPERTIES_SUM + GROUPS_NO + ENTITY_SUM
-
-SLOT_AND_GROUP = ENTITY_SUM + GROUPS_NO
-
-OUTPUT_GROUP_SLICE = np.s_[:, -GROUPS_NO:]
-OUTPUT_SLOT_SLICE = np.s_[:, -SLOT_AND_GROUP:-GROUPS_NO]
-OUTPUT_PROPERTIES_SLICE = np.s_[:, :-SLOT_AND_GROUP]
-
-OUTPUT_GROUP_SLICE_END = np.s_[-GROUPS_NO:]
-OUTPUT_SLOT_SLICE_END = np.s_[-SLOT_AND_GROUP:-GROUPS_NO]
-OUTPUT_PROPERTIES_SLICE_END = np.s_[:-SLOT_AND_GROUP]
-
-# Transformation
-# constant
-# progression -2, -1,1 ,2
-# arithmetic -/+ Position set arithmetic
-# distribute three
-
-# todo
-SLOTS_GROUPS = GROUPS_NO
-
-SLOT_TRANSFORMATION_NO = 4
-PROPERTY_TRANSFORMATION_NO = 8
-PROPERTIES_TRANSFORMATION_NO = PROPERTY_TRANSFORMATION_NO * PROPERTIES_NO
-PROPERTIES_TRANSFORMATION_SIZE = PROPERTIES_TRANSFORMATION_NO * ENTITY_SUM
-
-SLOT_TRANSFORMATION_SIZE = PROPERTY_TRANSFORMATION_NO * SLOTS_GROUPS
-INFERENCE_SIZE = SLOT_TRANSFORMATION_SIZE + PROPERTIES_TRANSFORMATION_SIZE
-
-INFERENCE_SLOT_SLICE = np.s_[:, :SLOT_TRANSFORMATION_SIZE]
-INFERENCE_PROPERTIES_SLICE = np.s_[:, -PROPERTIES_TRANSFORMATION_SIZE:]
-from operator import add
-
-
-# todo Refactor
-# Maybe properties should be on same level as rest.
-def decode_output(output, split_fn=np_split):
-    group_output = output[..., OUTPUT_GROUP_SLICE_END]
-    slot_output = output[..., OUTPUT_SLOT_SLICE_END]
-    properties_output = output[..., OUTPUT_PROPERTIES_SLICE_END]
-    properties_output_splited = split_fn(properties_output, list(rv.properties.INDEX.values()), axis=-1)
-    return group_output, slot_output, properties_output_splited
-
-
-def decode_inference(inference, reshape=np.reshape):
-    return reshape(inference[INFERENCE_SLOT_SLICE],
-                   [-1, SLOTS_GROUPS, PROPERTY_TRANSFORMATION_NO]), reshape(
-        inference[INFERENCE_PROPERTIES_SLICE],
-        [-1, PROPERTIES_NO, ENTITY_SUM, PROPERTY_TRANSFORMATION_NO])
-
-
-def decode_output_reshape(output, split_fn=np_split):
-    result = decode_output(output, split_fn=split_fn)
-    out_reshaped = []
-    for i, out in enumerate(result[2]):
-        shape = (-1, ENTITY_SUM, ENTITY_PROPERTIES_VALUES[i])
-        out_reshaped.append(out.reshape(shape))
-    return result[:2] + tuple(out_reshaped)
-
-
-def take_target(target):
-    return target[1], target[2]
-
-
-def create_target(images, index, pattern_index=(2, 5), full_index=False, arrange=np.arange, shape=lambda x: x.shape):
-    return [images[:, pattern_index[0]], images[:, pattern_index[1]],
-            images[arrange(shape(index)[0]), (0 if full_index else 8) + index[:, 0]]]
-
-
-def take_target_simple(target):
-    return target[1], target[0]
-
-
-def create_target_simple(images, target, index=slice(None), pattern_index=(2, 5)):
-    return [images[:, pattern_index[0]], images[:, pattern_index[1]], target][index]
-
-
-def decode_output_result(output, split_fn=np_split, arg_max=np.argmax):
-    result = decode_output_reshape(output, split_fn=split_fn)
-    res = []
-    for i, r in enumerate(result):
-        if i == 1:
-            res.append(r)
-        else:
-            res.append(arg_max(r, axis=-1))
-    return tuple(res)
-
-
-def decode_target(target):
-    target_group = target[..., 0]
-    target_slot = target[..., 1:INDEX[0]]
-    target_properties = target[..., INDEX[0]:END_INDEX]
-    target_properties_splited = [
-        target_properties[..., ::PROPERTIES_NO],
-        target_properties[..., 1::PROPERTIES_NO],
-        target_properties[..., 2::PROPERTIES_NO]
-    ]
-    return target_group, target_slot, target_properties_splited
-
-
-def decode_target_flat(target):
-    t = decode_target(target)
-    return t[0], t[1], t[2][0], t[2][1], t[2][2]
-
-
-def draw_board(images,  target=None, predict=None,image=None, desc=None, layout=None, break_=20):
-    if image != "target" and predict is not None:
-        image = images[predict:predict + 1]
-    elif images is None and target is not None:
-        image = images[target:target + 1]
-    # image = False to not draw anything
-    border = [{COR: target - 8, EXIST: (1, 3)}] + [{COR: p, EXIST: (0, 2)} for p in none(predict)]
-
-    boards = []
-    boards.append(draw_images(np.concatenate([images[:8], image[None] if len(image.shape)==3 else image]) if image is not None else images[:8]))
-    if layout == 1:
-        i = draw_images(images[8:], column=4, border=border)
-        if break_:
-            i = np.concatenate([np.zeros([ break_, i.shape[1],1]),i ],axis=0)
-        boards.append(i)
-
-    else:
-        boards.append(
-            draw_images(np.concatenate([images[8:], predict]) if predict is not None else images[8:], column=4,
-                        border=target - 8))
-    full_board = draw_images(boards, grid=False)
-    if desc:
-        full_board = add_text(full_board, desc)
-    return full_board
-
-
-def draw_boards(images, target=None, predict=None, image=None, desc=None, no=1, layout=None):
-    boards = []
-    for i, image in enumerate(images):
-        boards.append(draw_board(image, target[i][0] if target is not None else None,
-                                 predict[i] if predict is not None else None,
-                                 image[i] if image is not None else None,
-                                 desc[i] if desc is not None else None, layout=layout))
-    return boards
-
-
-def draw_raven(generator, predict=None, no=1, add_target_desc=True, indexes=None, types=TYPES,
-               layout=1):
-    if indexes is None:
-        indexes = val_sample(no)
-    data = generator.data[indexes]
-    if is_model(predict):
-        d = filter_keys(data, PROPERTY,reverse=True)
-        # tmp change
-        pro = predict(d)['predict']
-        print(pro)
-        predict = render_panels(pro, target=False)
-    # if target is not None:
-    target = data[TARGET]
-    target_index = data["index"]
-    images = data[INPUTS]
-
-    if hasattr(predict, "shape"):
-        if len(predict.shape) > 3:
-            # iamges
-            image = predict
-            # todo create index and output based on image
-            predict = None
-            predict_index = None
-        elif len(predict.shape) == 3:
-            image = render_panels(predict, target=False)
-            # Create index based on predict.
-            predict_index = None
-        else:
-            image = images[predict]
-            predict_index = predict
-            predict = target
-    else:
-        image = K.gather(images, target_index[:, 0])
-        predict_index = None
-        predict = None
-
-    # elif not(hasattr(target,"shape") and len(target.shape) > 3):
-    #     if hasattr(target,"shape") and target.shape[-1] == OUTPUT_SIZE:
-    #         pro = target
-    #         predict = render_panels(pro)
-    #     elif hasattr(target,"shape") and target.shape[-1] == FEATURE_NO:
-    #         # pro = target
-    #         pro = np.zeros([no, OUTPUT_SIZE], dtype="int")
-    #     else:
-    #         pro = np.zeros([no, OUTPUT_SIZE], dtype="int")
-    #         # predict = [None] * no
-    #         predict = render_panels(data[TARGET])
-
-    all_rules = []
-    for d in data[PROPERTY]:
-        rules = []
-        for j, rule_group in enumerate(d.findAll("Rule_Group")):
-            # rules_all.append(rule_group['id'])
-            for j, rule in enumerate(rule_group.findAll("Rule")):
-                rules.append(f"{rule['attr']} - {rule['name']}")
-            rules.append("")
-        all_rules.append(rules)
-    target_desc = get_desc(target)
-    if predict is not None:
-        predict_desc = decode_output_result(predict) if predict.shape[-1] == OUTPUT_SIZE else get_desc(predict)
-    else:
-        predict_desc = [None] * len(target_desc)
-    for a, po, to in zip(all_rules, predict_desc, target_desc):
-        # fl(predict_desc[-1])
-        if po is None:
-            po = [None] * len(to)
-        for p, t in zip(po, to):
-            a.extend(
-                [" ".join([str(i) for i in t])] + (
-                    [" ".join([str(i) for i in p]), ""] if p is not None else []
-                )
-            )
-        # a.extend([""] + [] + [""] + [" ".join(fl(p))])
-
-    # image = draw_boards(data[INPUTS],target=data["index"], predict=predict[:no], desc=all_rules, no=no,layer=layer)
-    image = draw_boards(images, target=target_index, predict=predict_index, image=image, desc=None, no=no,
-                        layout=layout)
-    return lu([(i, j) for i, j in zip(image, all_rules)])
-
-
-def val_sample(no=GROUPS_NO, base=3):
-    indexes = np.arange(no) * 2000 + base
-    return indexes
-
-
-def get_desc(target, exist=None, types=TYPES, sizes=SIZES):
-    decoded = decode_target(target)
-    exist = decoded[1] if exist is None else exist
-    taken = np.stack(take(decoded[2], np.array(exist, dtype=bool))).T
-
-    figures_no = np.sum(exist, axis=-1)
-    desc = np.split(taken, np.cumsum(figures_no))[:-1]
-    # figures_no = np.sum(exist, axis=-1)
-    # div = np.split(desc, np.cumsum(figures_no))[:-1]
-    result = []
-    for pd in desc:
-        r = []
-        for p in pd:
-            r.append([p[0], sizes[p[1]], types[p[2]]])
-        result.append(r)
-
-    return result
-
-
-# def get
-
-
-def get_description(inputs, predict, pro, no, types=TYPES, sizes=SIZES):
-    # target = inputs[1][2][:no]
-    target = inputs[TARGET]
-    target_group = target[:, 0]
-    target_exist = np.asarray(target[:, 1:ENTITY_SUM + 1], dtype="bool")
-    target_rest = target[:, ENTITY_SUM + 1:ENTITY_SUM + 1 + ENTITY_SUM * PROPERTIES_NO]
-    pro_reshaped = np.reshape(pro, (pro.shape[0], -1, PROPERTIES_NO))
-    target_reshaped = np.reshape(target_rest, (target_rest.shape[0], -1, PROPERTIES_NO))
-
-    # mask = np.repeat(target_exist, [4] * ENTITY_SUM, axis=-1)
-    # masked_result = np.repeat(target_exist, [4] * ENTITY_SUM, axis=-1)
-    pro_res = pro_reshaped[target_exist]
-    target_res = target_reshaped[target_exist]
-    figures_no = np.sum(target_exist, axis=-1)
-
-    pro_div = np.split(pro_res, np.cumsum(figures_no))[:-1]
-    target_div = np.split(target_res, np.cumsum(figures_no))[:-1]
-    pro_result_full = []
-    target_result_full = []
-    for pd, td in zip(pro_div, target_div):
-        pro_result = []
-        target_result = []
-        for p in pd:
-            pro_result.append([p[0], sizes[p[1]], types[p[2]]])
-        for t in td:
-            target_result.append([t[0], sizes[t[1]], types[t[2]]])
-        pro_result_full.append(pro_result)
-        target_result_full.append(target_result)
-
-    return pro_result_full, target_result_full
-
-
-def get_properties(target, types=TYPES, sizes=SIZES):
-    target_exist = np.asarray(target[:, 1:ENTITY_SUM + 1], dtype="bool")
-    target_rest = target[:, ENTITY_SUM + 1:ENTITY_SUM + 1 + ENTITY_SUM * PROPERTIES_NO]
-    target_reshaped = np.reshape(target_rest, (target_rest.shape[0], -1, PROPERTIES_NO))
-    target_res = target_reshaped[target_exist]
-    figures_no = np.sum(target_exist, axis=-1)
-    target_div = np.split(target_res, np.cumsum(figures_no))[:-1]
-    target_result_full = []
-    for td in target_div:
-        target_result = []
-        for t in td:
-            target_result.append([t[0], sizes[t[1]], types[t[2]]])
-        target_result_full.append(target_result)
-    return target_result_full
-
-
-def desc_properties(target, decode_fn=None, types=TYPES, sizes=SIZES):
-    if decode_fn is None:
-        if target.shape[1] == OUTPUT_SIZE:
-            decode_fn = decode_output_result
-        else:
-            decode_fn = decode_target
-
-    target_div = decode_fn(target)[2:]
-    target_result_full = []
-    for td in target_div:
-        target_result = []
-        for t in td:
-            target_result.append([t[0], sizes[t[1]], types[t[2]]])
-        target_result_full.append(target_result)
-    return target_result_full
-
-
-def get_pro(t, types=TYPES, sizes=SIZES):
-    return [int(t[0]), sizes[t[1]], types[t[2]]]
-
-
-def get_pro2(td, types=TYPES, sizes=SIZES):
-    target_result = []
-    for t in td:
-        target_result.append([int(t[0]), sizes[t[1]], types[t[2]]])
-    return target_result
-
-
-def get_pro3(target_div, types=TYPES, sizes=SIZES):
-    target_result_full = []
-    for td in target_div.to_list():
-        target_result = []
-        for t in td:
-            target_result.append([int(t[0]), sizes[t[1]], types[t[2]]])
-        target_result_full.append(target_result)
-    return target_result_full
-
-
-from models_utils import init_image as def_init_image, is_model
-
-init_image = partial(def_init_image, shape=(16, 8, 80, 80, 1))

raven_utils/models/attn.py

@@ -1,187 +0,0 @@
-from __future__ import print_function
-
-import tensorflow as tf
-from tensorflow.keras import backend as K
-from tensorflow.keras.layers import LSTMCell
-from tensorflow.keras.models import Model
-from tensorflow.keras.layers import Conv2D, Dense
-from tensorflow.keras.losses import mse
-from tensorflow.keras.models import clone_model
-from tensorflow.layers.base import InputSpec, Layer
-
-from models.dense import create_conv_model
-from models.utils import broadcast
-
-
-class ReflectionPadding2D(Layer):
-    def __init__(self, padding=(1, 1), **kwargs):
-        self.padding = tuple(padding)
-        self.input_spec = [InputSpec(ndim=4)]
-        super(ReflectionPadding2D, self).__init__(**kwargs)
-
-    def compute_output_shape(self, s):
-        """ If you are using "channels_last" configuration"""
-        return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])
-
-    def call(self, x, mask=None):
-        w_pad, h_pad = self.padding
-        return tf.pad(x, [[0, 0], [h_pad, h_pad], [w_pad, w_pad], [0, 0]], 'REFLECT')
-
-
-class Conv2Ref(Layer):
-    def __init__(self, padding=(1, 1), **kwargs):
-        self.padding = tuple(padding)
-        self.input_spec = [InputSpec(ndim=4)]
-        super(ReflectionPadding2D, self).__init__(**kwargs)
-
-    def compute_output_shape(self, s):
-        """ If you are using "channels_last" configuration"""
-        return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])
-
-    def call(self, x, mask=None):
-        w_pad, h_pad = self.padding
-        return tf.pad(x, [[0, 0], [h_pad, h_pad], [w_pad, w_pad], [0, 0]], 'REFLECT')
-
-
-class SegmentationNetwork(Model):
-
-    def __init__(self, filters=64, kernels=(3, 3)):
-        super(RecAE, self).__init__()
-        self.conv_1 = Conv2D(filters, kernels, padding=SAME)
-        self.conv_2 = Conv2D(filters, kernels, padding=SAME)
-
-    def call(self, inputs):
-        x = K.relu(inputs)
-        x = self.conv_1(x)
-        x = K.relu(x)
-        x = self.conv_2(x)
-        return x + inputs
-
-
-class QueryNetwork(Model):
-
-    def __init__(self, units=64):
-        super(RecAE, self).__init__()
-        self.conv_1 = Dense(units)
-        self.conv_2 = Dense(units)
-
-    def call(self, inputs):
-        x = K.relu(inputs)
-        x = self.conv_1(x)
-        x = K.relu(x)
-        x = self.conv_2(x)
-        return x + inputs
-
-
-class RecAE(Model):
-
-    def __init__(self, head, bottle, decoder):
-        super(RecAE, self).__init__()
-        self.head = head
-        self.bottle = bottle
-        self.base = clone_model(bottle)
-        self.decoder = decoder
-        self.segmentation_network = SegmentationNetwork()
-        self.query_network = QueryNetwork()
-        self.control = LSTMCell(64)
-        self.memory = LSTMCell(64)
-
-    def call(self, inputs):
-        feature = self.head(inputs)
-        segmentation = self.segmentation_network(feature)
-        control_base = self.base(feature)
-        h_c = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
-        h_m = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
-        shape = K.shape(feature)[:-1]
-        full_attention = tf.zeros(shape)[..., tf.newaxis]
-        full_image = tf.zeros(K.shape(inputs))
-        masks = []
-        ff = tf.zeros(K.shape(inputs))
-        scope = tf.ones(shape)[..., tf.newaxis]
-        for i in range(4):
-            r_c, h_c = self.control(tf.concat([control_base, h_m[0]], 1), h_c)
-            query = self.query_network(h_c[0])
-            log_attention = image_attention(segmentation, query)
-            attention = K.sigmoid(log_attention)
-            mask = attention * scope
-            scope = scope - mask
-            im = feature * mask
-            # im = feature
-            latent = self.bottle(im)
-            decoded = self.decoder(latent)
-            # self.add_loss(K.mean(-mse(full_attention, attention)))
-            # self.add_loss(K.mean(-mse(tf.ones(attention.shape), attention)))
-            full_attention += attention
-            big_mask = tf.image.resize(mask, K.shape(inputs)[1:-1])
-            ff += K.sigmoid(decoded)
-            full_image += K.sigmoid(decoded) * big_mask
-            r_m, h_m = self.memory(latent, h_m)
-            masks.append(big_mask)
-        self.add_loss(K.mean(mse(inputs, full_image)))
-        return full_image, masks
-
-
-# def image_attention(image, query, scale=True):
-@tf.function
-def image_attention(image, query):
-    log_attention = K.sum(query[:, tf.newaxis, tf.newaxis, :] * image, axis=-1, keepdims=True)
-    # if scale is not None:
-    log_attention /= tf.sqrt(tf.cast(K.shape(image)[-1], dtype=float))
-    return log_attention
-
-
-class RecAE_2(Model):
-
-    def __init__(self, head, bottle, decoder):
-        super(RecAE_2, self).__init__()
-        self.head = head
-        self.bottle = bottle
-        # self.base = clone_model(bottle)
-        self.base = self.bottle
-        self.decoder = decoder
-        self.segmentation_network = create_conv_model((64, 64, 1))
-        self.control = LSTMCell(64)
-        self.memory = LSTMCell(64)
-
-    def call(self, inputs):
-        feature = self.head(inputs)
-        control_base = self.base(feature)
-        h_c = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
-        h_m = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
-        shape = K.shape(feature)[:-1]
-        full_attention = tf.zeros(shape)[..., tf.newaxis]
-        full_image = tf.zeros(K.shape(inputs))
-        big_masks = []
-        masks = []
-        ff = tf.zeros(K.shape(inputs))
-        scope = tf.ones(shape)[..., tf.newaxis]
-        for i in range(4):
-            if i ==3:
-                mask = scope
-            else:
-                r_c, h_c = self.control(tf.concat([control_base, h_m[0]], 1), h_c)
-                query = broadcast(h_c[0], feature.shape[1:])
-                log_attention = self.segmentation_network(tf.concat([feature, query], axis=-1))
-                attention = K.sigmoid(log_attention)
-                mask = attention * scope
-                scope = scope - mask
-            masks.append(mask)
-            im = feature * mask
-            # im = feature
-            latent = self.bottle(im)
-            decoded = self.decoder(latent)
-            # self.add_loss(K.mean(-mse(scope, mask)))
-            sum = K.sum(tf.ones(K.shape(mask)))
-            self.add_loss(K.abs((sum/4)-K.sum(mask))/sum)
-            # self.add_loss(K.mean(-mse(tf.zeros(K.shape(mask)), mask)))
-            for m in masks:
-                self.add_loss(K.mean(-mse(m,mask)))
-
-            full_attention += mask
-            big_mask = tf.image.resize(mask, K.shape(inputs)[1:-1])
-            ff += K.sigmoid(decoded)
-            full_image += K.sigmoid(decoded) * big_mask
-            r_m, h_m = self.memory(latent, h_m)
-            big_masks.append(big_mask)
-        self.add_loss(K.mean(mse(inputs, full_image)))
-        return full_image, big_masks

raven_utils/models/attn2.py

@@ -1,187 +0,0 @@
-from __future__ import print_function
-
-import tensorflow as tf
-from tensorflow.keras import backend as K
-from tensorflow.keras.layers import LSTMCell
-from tensorflow.keras.models import Model
-from tensorflow.keras.layers import Conv2D, Dense
-from tensorflow.keras.losses import mse
-from tensorflow.keras.models import clone_model
-from tensorflow.layers.base import InputSpec, Layer
-
-from models.dense import create_conv_model
-from models.utils import broadcast
-
-
-class ReflectionPadding2D(Layer):
-    def __init__(self, padding=(1, 1), **kwargs):
-        self.padding = tuple(padding)
-        self.input_spec = [InputSpec(ndim=4)]
-        super(ReflectionPadding2D, self).__init__(**kwargs)
-
-    def compute_output_shape(self, s):
-        """ If you are using "channels_last" configuration"""
-        return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])
-
-    def call(self, x, mask=None):
-        w_pad, h_pad = self.padding
-        return tf.pad(x, [[0, 0], [h_pad, h_pad], [w_pad, w_pad], [0, 0]], 'REFLECT')
-
-
-class Conv2Ref(Layer):
-    def __init__(self, padding=(1, 1), **kwargs):
-        self.padding = tuple(padding)
-        self.input_spec = [InputSpec(ndim=4)]
-        super(ReflectionPadding2D, self).__init__(**kwargs)
-
-    def compute_output_shape(self, s):
-        """ If you are using "channels_last" configuration"""
-        return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])
-
-    def call(self, x, mask=None):
-        w_pad, h_pad = self.padding
-        return tf.pad(x, [[0, 0], [h_pad, h_pad], [w_pad, w_pad], [0, 0]], 'REFLECT')
-
-
-class SegmentationNetwork(Model):
-
-    def __init__(self, filters=64, kernels=(3, 3)):
-        super(RecAE, self).__init__()
-        self.conv_1 = Conv2D(filters, kernels)
-        self.conv_2 = Conv2D(filters, kernels)
-
-    def call(self, inputs):
-        x = K.relu(inputs)
-        x = self.conv_1(x)
-        x = K.relu(x)
-        x = self.conv_2(x)
-        return x + inputs
-
-
-class QueryNetwork(Model):
-
-    def __init__(self, units=64):
-        super(RecAE, self).__init__()
-        self.conv_1 = Dense(units)
-        self.conv_2 = Dense(units)
-
-    def call(self, inputs):
-        x = K.relu(inputs)
-        x = self.conv_1(x)
-        x = K.relu(x)
-        x = self.conv_2(x)
-        return x + inputs
-
-
-class RecAE(Model):
-
-    def __init__(self, head, bottle, decoder):
-        super(RecAE, self).__init__()
-        self.head = head
-        self.bottle = bottle
-        self.base = clone_model(bottle)
-        self.decoder = decoder
-        self.segmentation_network = SegmentationNetwork()
-        self.query_network = QueryNetwork()
-        self.control = LSTMCell(64)
-        self.memory = LSTMCell(64)
-
-    def call(self, inputs):
-        feature = self.head(inputs)
-        segmentation = self.segmentation_network(feature)
-        control_base = self.base(feature)
-        h_c = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
-        h_m = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
-        shape = K.shape(feature)[:-1]
-        full_attention = tf.zeros(shape)[..., tf.newaxis]
-        full_image = tf.zeros(K.shape(inputs))
-        masks = []
-        ff = tf.zeros(K.shape(inputs))
-        scope = tf.ones(shape)[..., tf.newaxis]
-        for i in range(10):
-            r_c, h_c = self.control(tf.concat([control_base, h_m[0]], 1), h_c)
-            query = self.query_network(h_c[0])
-            log_attention = image_attention(segmentation, query)
-            attention = K.softmax(log_attention)
-            mask = attention * scope
-            scope = scope - mask
-            im = feature * mask
-            # im = feature
-            latent = self.bottle(im)
-            decoded = self.decoder(latent)
-            # self.add_loss(K.mean(-mse(full_attention, attention)))
-            # self.add_loss(K.mean(-mse(tf.ones(attention.shape), attention)))
-            full_attention += attention
-            big_mask = tf.image.resize(mask, K.shape(inputs)[1:-1])
-            ff += K.sigmoid(decoded)
-            full_image += K.sigmoid(decoded) * big_mask
-            r_m, h_m = self.memory(latent, h_m)
-            masks.append(big_mask)
-        self.add_loss(K.mean(mse(inputs, full_image)))
-        return full_image, masks
-
-
-# def image_attention(image, query, scale=True):
-@tf.function
-def image_attention(image, query):
-    log_attention = K.sum(query[:, tf.newaxis, tf.newaxis, :] * image, axis=-1, keepdims=True)
-    # if scale is not None:
-    log_attention /= tf.sqrt(tf.cast(K.shape(image)[-1], dtype=float))
-    return log_attention
-
-
-class RecAE_2(Model):
-
-    def __init__(self, head, bottle, decoder):
-        super(RecAE_2, self).__init__()
-        self.head = head
-        self.bottle = bottle
-        # self.base = clone_model(bottle)
-        self.base = self.bottle
-        self.decoder = decoder
-        self.segmentation_network = create_conv_model((64, 64, 1))
-        self.control = LSTMCell(64)
-        self.memory = LSTMCell(64)
-
-    def call(self, inputs):
-        feature = self.head(inputs)
-        control_base = self.base(feature)
-        h_c = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
-        h_m = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
-        shape = K.shape(feature)[:-1]
-        full_attention = tf.zeros(shape)[..., tf.newaxis]
-        full_image = tf.zeros(K.shape(inputs))
-        big_masks = []
-        masks = []
-        ff = tf.zeros(K.shape(inputs))
-        scope = tf.ones(shape)[..., tf.newaxis]
-        for i in range(4):
-            if i ==3:
-                mask = scope
-            else:
-                r_c, h_c = self.control(tf.concat([control_base, h_m[0]], 1), h_c)
-                query = broadcast(h_c[0], feature.shape[1:])
-                log_attention = self.segmentation_network(tf.concat([feature, query], axis=-1))
-                attention = K.sigmoid(log_attention)
-                mask = attention * scope
-                scope = scope - mask
-            masks.append(mask)
-            im = feature * mask
-            # im = feature
-            latent = self.bottle(im)
-            decoded = self.decoder(latent)
-            # self.add_loss(K.mean(-mse(scope, mask)))
-            sum = K.sum(tf.ones(K.shape(mask)))
-            self.add_loss(K.abs((sum/4)-K.sum(mask))/sum)
-            # self.add_loss(K.mean(-mse(tf.zeros(K.shape(mask)), mask)))
-            for m in masks:
-                self.add_loss(K.mean(-mse(m,mask)))
-
-            full_attention += mask
-            big_mask = tf.image.resize(mask, K.shape(inputs)[1:-1])
-            ff += K.sigmoid(decoded)
-            full_image += K.sigmoid(decoded) * big_mask
-            r_m, h_m = self.memory(latent, h_m)
-            big_masks.append(big_mask)
-        self.add_loss(K.mean(mse(inputs, full_image)))
-        return full_image, big_masks

raven_utils/models/body.py

@@ -1,276 +0,0 @@
-import itertools
-
-import tensorflow as tf
-from ml_utils import self_product, lw
-
-from models_utils import DictModel, ListModel, Flat, bm, Base, Cat, Res, Flat2, conv, KERNEL_SIZE, FILTERS, SAME, \
-    Get, SM, bs, RELU, ACTIVATION, dense, bd, HardBlock, MaxBlock
-import models_utils.ops as K
-from models_utils import Merge, SoftBlock
-from models_utils.build import build_multi_dense, build_multi_conv, build_conv_model, build_encoder
-from tensorflow.keras.layers import Lambda, Dense
-from tensorflow.keras.layers import Conv2D
-
-from config.constant import MEMORY, CONTROL, LATENT, MERGE, CONCAT, INFERENCE, FLATTEN
-from models_utils.config import config
-
-
-class RavRes(Res):
-    def __init__(self, model="v2", latent=256, act=RELU):
-        super().__init__(model=model)
-        self.latent = latent
-
-    def call(self, inputs):
-        return self.model(inputs) + inputs[0][:, ..., self.latent:]
-
-
-# not working
-class RavResConv(Res):
-    def __init__(self, model="v2", latent=256, act=RELU):
-        super().__init__(model=model)
-        self.latent = latent
-        self.conv = conv(latent, (1, 1), activation=act)
-
-    def call(self, inputs):
-        return self.model(inputs) + self.conv(inputs[0])
-
-
-class RavResDense(Res):
-    def __init__(self, model="v2", latent=256, act=config.DEF_DENSE.activation):
-        super().__init__(model=model)
-        self.latent = latent
-        self.conv = dense(latent, activation=act)
-
-    def call(self, inputs):
-        return self.model(inputs) + self.conv(inputs[0])
-
-
-def create_dense_block(latent=256, loop=1):
-    soft_block = Res(SoftBlock(build_multi_dense(latent), add_identity=None,
-                               score_activation=tf.sigmoid), latent=latent)
-    cells = [
-        (lambda x: K.cat([x[:, 0], x[:, 1]]), LATENT, CONCAT),
-        (None, CONCAT, MEMORY),
-        (Dense(latent), CONCAT, MERGE),
-        (Merge(latent), [INFERENCE, MERGE], CONTROL),
-        (soft_block, [MEMORY, CONTROL], MEMORY)
-    ]
-
-    return ListModel([DictModel(*cell) for cell in cells] * loop, [LATENT, INFERENCE], MEMORY)
-
-
-def build_multi_conv(filters=32, end_filters=64, padding="same",mul=1, norm=None, **kwargs):
-    base = [(1, 3), (3, 1), (3, 3)]
-    block = list(self_product(base))
-    block2 = [b + b[0:1] for b in block]
-    block3 = [b + b for b in block]
-    block4 = ([[(3, 3)]] + [[(3, 3), (3, 3)]] + [[(3, 3), (3, 3), (3, 3)]]) * 2
-    block5 = [[], []]
-    all_blocks = [s for b in [block, block2, block3, block4, block5] for s in b]
-    start = {
-        FILTERS: filters,
-        KERNEL_SIZE: (1, 1)
-    }
-
-    end = {
-        FILTERS: end_filters,
-        KERNEL_SIZE: (1, 1),
-        ACTIVATION: None
-    }
-
-    all_arch = []
-    for ab in all_blocks:
-        arch = [{
-            FILTERS: filters,
-            KERNEL_SIZE: a,
-            **kwargs
-        } for a in ab]
-        all_arch.append([start] + arch + [end])
-
-    all_arch = all_arch * mul
-
-    return [
-        build_encoder(a, add_norm=norm if norm else None, padding=padding, name=f"b{i}", order=(1, 0) if norm else None)
-        for i, a in enumerate(all_arch)]
-
-
-def create_block(latent=256, simpler=0, loop=1, padding=SAME, norm=None, trans_div=2, act="pass", type_="conv",
-                 block_=SoftBlock,max_k=16,
-                 **kwargs):
-    trans_size = int(latent / trans_div)
-    # if block_ == HardBlock:
-    #     mul = 2
-    # elif block_ == MaxBlock:
-    #     mul = int(38/max_k)
-    # else:
-    #     mul = 1
-
-    if act == "pass":
-        res_class = RavRes
-    else:
-        if type_ == "dense":
-            res_class = RavResDense
-        else:
-            res_class = RavResConv
-
-    if type_ == "dense":
-        build_res = lambda: Res(model="dv2")
-        # build_reduction = lambda: bm([dense(latent), "IN"])
-        build_reduction = lambda: dense(latent)
-        build_flatten = lambda: bd([latent] * 2)
-    else:
-        build_res = lambda: Res(padding=padding)
-        build_reduction = lambda: bm([conv(trans_size if simpler else latent, 1, padding=padding), "BN"])
-        # build_reduction = lambda: bm([conv(latent, 1, padding=padding), "BN"])
-        # build_reduction = lambda: bm([conv(trans_size, 1, padding=padding), "BN"])
-        # build_reduction = lambda: conv(trans_size, 1, padding=padding)
-        # build_flatten = lambda: Flat2(filters=trans_size,res_no=2, padding=padding, units=64)
-        build_flatten = lambda: Flat2(filters=trans_size,padding=padding, units=64)
-
-    if simpler == 1:
-        cells = [
-            (lambda x: K.cat([x[:, 0], x[:, 1]]), LATENT, CONCAT,"concatenation"),
-            # (None, CONCAT, MEMORY),
-            (build_reduction(), CONCAT, MERGE,"Start_resnet_block"),
-            # (Get(), INFERENCE, INFERENCE),
-            (K.cat, [INFERENCE, MERGE], CONTROL,"concatenation"),
-        ]
-    else:
-        cells = [
-            (lambda x: K.cat([x[:, 0], x[:, 1]]), LATENT, CONCAT),
-            (build_reduction(), CONCAT, MEMORY),
-            (build_reduction(), INFERENCE, CONTROL),
-        ]
-    for i, l in enumerate(lw(loop)):
-        if l:
-            concat = K.cat
-            control_reduction = build_reduction()
-            control_res = build_res()
-            control_flatten = build_flatten()
-            if i == 0 and simpler == 1:
-                rest_params = {
-                    "latent": latent,
-                    "act": act
-                }
-            else:
-                rest_params = {
-                    "latent": 0
-                }
-
-
-            if block_ == SoftBlock:
-                block_params = {
-                }
-            else:
-                block_params = {
-                    "trans_output_shape": latent
-                }
-                if block_ == MaxBlock:
-                    block_params['max_k'] = max_k
-
-
-            # todo change name
-            soft_block = res_class(
-                block_(
-                    build_multi_dense(latent) if type_ == "dense" else build_multi_conv(trans_size, end_filters=latent,
-                                                                                        norm=norm, padding=padding,
-                                                                                        **kwargs),
-                    add_identity=None,
-                    score_activation=tf.sigmoid,
-                    **block_params
-
-                ),
-                **rest_params)
-
-        if i == 0 and simpler == 1:
-            cells.extend([
-                (control_reduction, CONTROL, CONTROL,"Reduction"),
-                (control_res, CONTROL, CONTROL,"Control_resnet_block"),
-                (control_flatten, CONTROL, FLATTEN,"Weights"),
-                (soft_block, [CONCAT, FLATTEN], MEMORY,"Transformation"),
-                # (soft_block, [MEMORY, FLATTEN], MEMORY,"Transformation"),
-            ])
-        else:
-            if l:
-                memory_res = build_res()
-
-            cells.extend([
-                (memory_res, MEMORY, MEMORY,"Memory_resnet_block"),
-                (concat, [CONTROL, MEMORY], CONTROL,"concatenation"),
-                (control_reduction, CONTROL, CONTROL,"Reduction"),
-                (control_res, CONTROL, CONTROL,"Control_resnet_block"),
-                (control_flatten, CONTROL, FLATTEN,"Weights"),
-                (soft_block, [MEMORY, FLATTEN], MEMORY, "Transformation"),
-            ])
-    return ListModel([DictModel(*cell) for cell in cells], [LATENT, INFERENCE], MEMORY, debug_=False)
-
-#
-#
-# def test(x):
-#     np.zeros(4)
-#     self_product((1, 3))
-#
-#
-# list(itertools.product())
-# u.layers[0].layers[-1].model.layers[1]
-
-# class RecurrentBodyDict(Model):
-#     # def __init__(self, start=None, cell=None, output_network=None, output_activation="tanh", latent=64, loop_no=5):
-#     def __init__(self, start=None, cell=None, output_network=None, output_activation=None, latent=64, loop_no=5):
-#         super().__init__()
-#         self.start = sm(start, lambda: SubClassingModel([StartLSTMControl(latent), StartLSTMMemory(latent)]),
-#                         latent=latent)
-#         self.cell = sm(cell, lambda: SubClassingModel([LSTMControl(latent), LSTMMemory(latent)]), latent=latent)
-#         self.output_network = sm(output_network, lf(take_memory_states))
-#         self.loop_no = loop_no
-#         # tmp
-#         self.activation = Activation(output_activation)
-#
-#     def call(self, inputs):
-#         outputs = []
-#         for j in range(3):
-#             outputs.append(self.start({"latent": inputs[0][j], "inference": inputs[1]}))
-#         for i in range(self.loop_no):
-#             for j in range(3):
-#                 outputs[j] = self.cell(outputs[j])
-#
-#         return self.activation(self.output_network(outputs))
-#
-#
-# class RecurrentBodySimpleMix4Dict(RecurrentBodyDict):
-#     def __init__(self, latent=64, output_network=None, loop_no=5):
-#         super().__init__(
-#             start=SubClassingModel(
-#                 [ConcatCell(), DenseCell(latent), InfMergeCell(latent),
-#                  WeigthCell(latent, layer_no=np.repeat([1, 2, 3, 4, 5, 6, 7, 8], 4),
-#                             add_identity=Lambda(lambda x: x[:, latent:]))]),
-#             cell=False,
-#             output_network=output_network, loop_no=0)
-# class RecurrentBodySimpleMix4Conv(RecurrentBodyDict):
-#     def __init__(self, latent=64, output_network=None, loop_no=5):
-#         super().__init__(
-#             start=SubClassingModel(
-#                 [ConcatCell(), ConvCell(latent), ReduceCell(latent), InfMergeCell(latent),
-#                  ModelCell(latent=latent, layers_no=2, input_name=CONTROL, result_name=CONTROL),
-#                  WeigthCell(latent,
-#                             transformation_network=[build_conv_model2([latent] * i, kernels=(j, j)) for i in range(1, 7) for j in
-#                                                     range(1, 5) for _ in range(1)],
-#                             add_identity=Lambda(lambda x: x[:, ..., latent:]))
-#                  ]),
-#             cell=False,
-#             output_network=output_network, loop_no=0)
-#
-#
-# class RecurrentBodySimpleMix4Conv2(RecurrentBodyDict):
-#     def __init__(self, latent=64, output_network=None, loop_no=5):
-#         super().__init__(
-#             start=SubClassingModel(
-#                 [ConcatCell(), ConvCell(latent), ReduceCell2(latent), InfMergeCell(latent),
-#                  ModelCell(latent=latent, layers_no=2, input_name=CONTROL, result_name=CONTROL),
-#                  WeigthCell(latent,
-#                             transformation_network=[bc([latent] * i, kernels=(j, j)) for i in range(1, 7) for j in
-#                                                     range(1, 5) for _ in range(1)],
-#                             add_identity=Lambda(lambda x: x[:, ..., latent:]))
-#                  ]),
-#             cell=False,
-#             output_network=output_network, loop_no=0)

raven_utils/models/class_.py

@@ -1,31 +0,0 @@
-from ml_utils import lw
-from models_utils import SubClassingModel, ops as K, Base
-import tensorflow as tf
-
-
-class Merge(SubClassingModel):
-    def call(self, inputs):
-        results = []
-        for i, model in enumerate(self.model[:-1]):
-            results.append(model(inputs[i]))
-        # todo why K.cat not working
-        results = self.model[-1](tf.concat(results, axis=-1))
-        return results
-
-
-class RavenClass(Base):
-    def __init__(self, model, scales=None, no=3, name=None):
-        super().__init__(model=model, name=name)
-        self.scales = scales
-        self.no = no
-
-    def call(self, inputs):
-        inputs = lw(inputs)
-        class_res = []
-        # for i in range(inputs[0].shape[1]):
-        for i in range(self.no):
-            # d = [r[:, i] if r.ndim == 5 else r for r in inputs]
-            d = [inputs[s][:, i] if inputs[s].ndim > 2 else inputs for s in self.scales]
-            class_res.append(self.model(d))
-        # return tf.stack(class_res,axis=1)
-        return [class_res]

raven_utils/models/head.py

@@ -1,159 +0,0 @@
-import tensorflow as tf
-from ml_utils import set_default
-from models_utils import build_dense_model, bm, ActivationModel, sm, large_conv_dense_encoder, Pass
-from models_utils import res
-from tensorflow.keras import Model
-from models_utils import ops as K
-from tensorflow.keras.layers import Dense, Conv2D, Flatten
-from keras.backend import batch_flatten
-
-
-# todo Refactoring
-class HeadModel(Model):
-    def __init__(self, encoder=None, inference_network=None, output_size=64, inference_output_size=None,
-                 inference_activation="relu", stem=None, images_no=8, inference_image_no=None):
-        super().__init__()
-        # self.encoder = sm(encoder, bm([en.large_conv_dense_encoder(), Dense(output_size)], False))
-        self.encoder = encoder or bm([large_conv_dense_encoder(), Dense(output_size)])
-        # self.head = head or HeadBatch(encoder=encoder, output_size=output_size)
-        inference_output_size = inference_output_size or output_size
-        self.inference_network = inference_network or bm([
-            K.flat,
-            build_dense_model([1028, 512, 512, inference_output_size],
-                              last_activation=inference_activation)]
-        )
-        self.stem = stem or Pass()
-        self.images_no = images_no
-        self.inference_image_no = self.images_no if inference_image_no is None else inference_image_no
-
-
-class LatentHeadModel(HeadModel):
-    def call(self, inputs):
-        result = K.map_batch(inputs[:, :self.images_no], self.encoder)
-        inference = self.inference_network(result[:, :self.inference_image_no])
-        latents = self.stem(result)
-        return [latents, inference,result]
-
-
-# # todo use map_batch
-# class HeadBatch(Model):
-#     def __init__(self, encoder=None, output_size=64):
-#         super().__init__()
-#         self.encoder = sm(encoder, bm([large_conv_dense_encoder(), Dense(output_size)], False))
-#
-#     def call(self, inputs):
-#         shape = tf.shape(inputs)
-#         latents = self.encoder(tf.reshape(inputs, shape=tf.concat([[-1], shape[2:]], axis=-1)))
-#         latents = K.reshape(latents, tf.concat([[-1, shape[1]], latents.shape[1:]], axis=-1))
-#         return latents
-
-
-# Not working
-class DuoHeadModel(HeadModel):
-    def __init__(self, encoder=None, inference_network=None, images_no=8, filters=-4):
-        super().__init__(encoder=encoder, inference_network=inference_network, images_no=images_no)
-        self.encoder = ActivationModel(self.encoder, filters=filters, include_input=False)
-
-    def call(self, inputs):
-        shape = inputs.shape
-        result = reversed(self.encoder(K.reshape(inputs, shape=[-1] + list(shape[2:]))))
-        latents = K.reshape(result[0], [-1, self.images_no] + [result[0].shape[-1]])
-        inference = self.inference_network(K.flat(result[1]))
-        return [latents, inference]
-
-
-class MultiHeadModel(Model):
-    def __init__(self, encoder=None, images_no=8, filters=(1, 3, 6)):
-        super().__init__()
-        self.encoder = ActivationModel(encoder, filters=filters, include_input=False)
-        self.merge = MergeSacles()
-        self.images_no = images_no
-
-    def call(self, inputs):
-        shape = tf.shape(inputs)
-        results = self.encoder(tf.reshape(inputs, shape=tf.concat([[-1], shape[2:]], axis=-1)))
-        latents = [tf.reshape(result, shape=tf.concat([[-1, self.images_no], tf.shape(result)[1:]], axis=-1)) for result
-                   in results]
-
-        l1 = tf.transpose(latents[0], (0, 2, 3, 1, 4))
-        # l1 = tf.reshape(l1, tuple(list(l1.shape[:3]) + [l1.shape[-2] * l1.shape[-1]]))
-        shape = tf.shape(l1)
-        l1 = tf.reshape(l1, tf.concat([[-1], shape[1:3], [shape[-2] * shape[-1]]], axis=-1))
-
-        l2 = tf.transpose(latents[1], (0, 2, 3, 1, 4))
-        # l2 = tf.reshape(l2, [-1] + list(l2.shape[1:3]) + [l2.shape[-2] * l2.shape[-1]])
-        shape = tf.shape(l2)
-        l2 = tf.reshape(l2, tf.concat([[-1], shape[1:3], [shape[-2] * shape[-1]]], axis=-1))
-
-        l3 = latents[2]
-        shape = tf.shape(l3)
-        # l3 = tf.reshape(l3, [-1] + [l3.shape[-2] * l3.shape[-1]])
-        l3 = tf.reshape(l3, tf.concat([[-1], [shape[-2] * shape[-1]]], axis=-1))
-
-        inference = self.merge([l1, l2, l3])
-        return [latents, inference]
-
-
-class MergeSacles(Model):
-    def __init__(self):
-        super().__init__()
-        self.inf_1 = bm([Conv2D(64, 1, activation="relu"), res(64),
-                         Conv2D(64, 3, strides=2, padding=SAME, activation="relu"),
-                         res(64),
-                         Flatten(),
-                         Dense(256, "relu")])
-        self.inf_2 = bm([Conv2D(128, 1, activation="relu"),
-                         res(128),
-                         Flatten(),
-                         Dense(256, "relu")])
-        self.inf_3 = Dense(256, "relu")
-
-    def call(self, inputs):
-        il1 = self.inf_1(inputs[0])
-        il2 = self.inf_2(inputs[1])
-        il3 = self.inf_3(inputs[2])
-        inference = tf.concat([il1, il2, il3], axis=1)
-        return inference
-
-
-class MultiHeadModel2(Model):
-    def __init__(self, encoder=None, images_no=8, filters=(3, 6)):
-        super().__init__()
-        self.encoder = ActivationModel(encoder, filters=filters, include_input=False)
-        self.merge = MergeSacles2()
-        self.images_no = images_no
-
-    def call(self, inputs):
-        shape = tf.shape(inputs)
-        results = self.encoder(tf.reshape(inputs, shape=tf.concat([[-1], shape[2:]], axis=-1)))
-        latents = [tf.reshape(result, shape=tf.concat([[-1, self.images_no], tf.shape(result)[1:]], axis=-1)) for result
-                   in results]
-
-        l2 = tf.transpose(latents[0], (0, 2, 3, 1, 4))
-        # l2 = tf.reshape(l2, [-1] + list(l2.shape[1:3]) + [l2.shape[-2] * l2.shape[-1]])
-        shape = tf.shape(l2)
-        l2 = tf.reshape(l2, tf.concat([[-1], shape[1:3], [shape[-2] * shape[-1]]], axis=-1))
-
-        l3 = latents[1]
-        shape = tf.shape(l3)
-        # l3 = tf.reshape(l3, [-1] + [l3.shape[-2] * l3.shape[-1]])
-        l3 = tf.reshape(l3, tf.concat([[-1], [shape[-2] * shape[-1]]], axis=-1))
-
-        inference = self.merge([l2, l3])
-        return [latents, inference]
-
-
-class MergeSacles2(Model):
-    def __init__(self):
-        super().__init__()
-        self.inf_1 = bm([Conv2D(128, 1, activation="relu"),
-                         res(128),
-                         Flatten(),
-                         Dense(256, "relu")])
-        self.inf_2 = Dense(256, "relu")
-
-    def call(self, inputs):
-        il1 = self.inf_1(inputs[0])
-        il2 = self.inf_2(inputs[1])
-        inference = tf.concat([il1, il2], axis=1)
-        return inference

raven_utils/models/loss.py

@@ -1,630 +0,0 @@
-from functools import partial
-
-import tensorflow as tf
-import tensorflow.experimental.numpy as tnp
-from models_utils import OUTPUT, TARGET, PREDICT, DictModel, add_loss, LOSS, Predict
-from models_utils import SubClassingModel
-from models_utils.models.utils import interleave
-from models_utils.op import reshape
-from tensorflow.keras import Model
-# from tensorflow.keras import backend as K
-from tensorflow.keras.layers import Lambda
-from tensorflow.keras.losses import SparseCategoricalCrossentropy, mse
-from tensorflow.keras.metrics import SparseCategoricalAccuracy, Accuracy, BinaryAccuracy
-import models_utils.ops as K
-
-import raven_utils.decode
-import raven_utils as rv
-from raven_utils.config.constant import LABELS, INDEX, ACC_SAME, ACC_CHOOSE_LOWER, ACC_CHOOSE_UPPER, CLASSIFICATION, \
-    SLOT, \
-    PROPERTIES, ACC, GROUP, NUMBER, MASK
-from raven_utils.models.uitls_ import RangeMask
-from raven_utils.const import VERTICAL, HORIZONTAL
-
-
-def get_properties_mask(target):
-    return target[:, rv.target.END_INDEX:rv.target.UNIFORMITY_INDEX] > 0
-
-
-def create_change_mask(target):
-    properties_mask = get_properties_mask(target)
-    return [create_mask(properties_mask, i) for i, _ in enumerate(rv.rules.ATTRIBUTES)]
-
-
-def create_uniform_mask(target):
-    u_mask = lambda i: tf.tile(target[:, rv.target.UNIFORMITY_INDEX + i, None] == 3, [1, rv.rules.ATTRIBUTES_LEN])
-    properties_mask = tf.concat([u_mask(0), u_mask(1)], axis=-1) | get_properties_mask(target)
-    return [create_mask(properties_mask, i) for i, _ in enumerate(rv.rules.ATTRIBUTES)]
-
-
-def create_all_mask(target):
-    return [
-        tf.cast(tf.ones(tf.stack([tf.shape(target)[0], rv.entity.SUM])), dtype=tf.bool) for i, _ in
-        enumerate(rv.rules.ATTRIBUTES)]
-
-
-class BaselineClassificationLossModel(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True):
-        super().__init__()
-        self.predict_fn = SubClassingModel([lambda x: x[0], PredictModel()])
-        self.loss_fn = ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                       group_loss=group_loss)
-        self.metric_fn = SimilarityRaven(mode=mode)
-
-    def call(self, inputs):
-        losses = []
-        output = inputs[1]
-        losses.append(self.loss_fn([inputs[0][0], output]))
-        losses.append(self.metric_fn([inputs[0][2], inputs[3][0], inputs[0][1][:, 8:]]))
-        return losses
-
-
-class RavenLoss(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.3),
-                 classification=False, trans=True, anneal=False):
-        super().__init__()
-        if anneal:
-            self.weight_scheduler
-        self.classification = classification
-        self.trans = trans
-        self.predict_fn = DictModel(SubClassingModel([lambda x: x[-1], PredictModel()]), in_=OUTPUT,
-                                    out=[PREDICT, MASK], name="pred")
-        if self.trans:
-            self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                                    group_loss=group_loss, enable_metrics=False, lw=lw[0]),
-                                    name="main_loss")
-            self.loss_fn_2 = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                                      group_loss=group_loss), name="add_loss")
-            self.metric_fn = SimilarityRaven(mode=mode)
-        if self.classification:
-            self.loss_fn_3 = add_loss(
-                ClassRavenModel(mode=create_all_mask, number_loss=number_loss, slot_loss=slot_loss,
-                                group_loss=group_loss, enable_metrics="c" if self.trans else True), lw=lw[1],
-                name="class_loss")
-
-    def call(self, inputs):
-        losses = []
-        output = inputs[OUTPUT]
-        target = inputs[TARGET]
-        labels = inputs[LABELS]
-
-        if self.trans:
-            losses.append(self.loss_fn([labels[:, 2], output[0]]))
-            losses.append(self.loss_fn([labels[:, 5], output[1]]))
-            losses.append(self.loss_fn_2([target, output[2]]))
-            losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
-        if self.classification:
-            for i in range(8):
-                losses.append(self.loss_fn_3([labels[:, i], inputs[CLASSIFICATION][i]]))
-        return {**inputs, LOSS: losses}
-
-
-class VTRavenLoss(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.1),
-                 classification=False, trans=True, anneal=False, plw=None):
-        super().__init__()
-        if anneal:
-            self.weight_scheduler
-        self.classification = classification
-        self.trans = trans
-        self.predict_fn = DictModel(SubClassingModel([lambda x: x[:, -1], PredictModel()]), in_=OUTPUT,
-                                    out=[PREDICT, MASK], name="pred")
-        self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                                group_loss=group_loss, plw=plw), lw=lw[0] , name="add_loss")
-        self.metric_fn = SimilarityRaven(mode=mode)
-        if self.classification:
-            self.loss_fn_2 = add_loss(
-                ClassRavenModel(mode=create_all_mask, number_loss=number_loss, slot_loss=slot_loss,
-                                group_loss=group_loss, enable_metrics="c", plw=plw), lw=lw[1],  name="class_loss")
-
-    def call(self, inputs):
-        losses = []
-        output = inputs[OUTPUT]
-        target = inputs[TARGET]
-        labels = inputs[LABELS]
-
-        for i in range(9):
-            losses.append(self.loss_fn_2([labels[:, i], output[:, i]]))
-        losses.append(self.loss_fn([target, output[:, 8]]))
-        losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
-        return {**inputs, LOSS: losses}
-
-
-class SingleVTRavenLoss(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.1),
-                 classification=False, trans=True, anneal=False):
-        super().__init__()
-        if anneal:
-            self.weight_scheduler
-        self.classification = classification
-        self.trans = trans
-        self.predict_fn = DictModel(PredictModel(), in_=OUTPUT, out=[PREDICT, MASK], name="pred")
-        self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                                group_loss=group_loss), lw=lw[0], name="add_loss")
-        self.metric_fn = SimilarityRaven(mode=mode)
-
-    def call(self, inputs):
-        losses = []
-        output = inputs[OUTPUT]
-        target = inputs[TARGET]
-        labels = inputs[LABELS]
-
-        losses.append(self.loss_fn([target, output]))
-        losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
-        return {**inputs, LOSS: losses}
-
-
-class ClassRavenModel(Model):
-    def __init__(self, mode=create_all_mask,plw=None, number_loss=False, slot_loss=True, group_loss=True, enable_metrics=True,
-                 lw=1.0):
-        super().__init__()
-        self.number_loss = number_loss
-        self.group_loss = group_loss
-        self.enable_metrics = enable_metrics
-        self.slot_loss = slot_loss
-        self.predict_fn = PredictModel()
-        self.loss_fn = SparseCategoricalCrossentropy(from_logits=True)
-        if self.slot_loss:
-            self.loss_fn_2 = tf.nn.sigmoid_cross_entropy_with_logits
-        if self.enable_metrics:
-            self.enable_metrics = f"{self.enable_metrics}_" if isinstance(self.enable_metrics, str) else ""
-            self.metric_fn = [
-                SparseCategoricalAccuracy(name=f"{self.enable_metrics}{ACC}_{property_}") for property_ in
-                rv.properties.NAMES]
-            if self.group_loss:
-                self.metric_fn_group = SparseCategoricalAccuracy(name=f"{self.enable_metrics}{ACC}_{GROUP}")
-            if self.slot_loss:
-                self.metric_fn_2 = BinaryAccuracy(name=f"{self.enable_metrics}{ACC}_{SLOT}")
-        self.range_mask = RangeMask()
-        self.mode = mode
-        self.lw = lw
-        if not plw:
-            plw = [1., 95.37352927, 2.83426987, 0.85212836, 1.096005, 1.21943385]
-        elif isinstance(plw, int) or isinstance(plw, float):
-            plw = [1., plw, 2.83426987, 0.85212836, 1.096005, 1.21943385]
-            # plw = [plw] * 6
-        self.plw = plw
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    def call(self, inputs):
-        losses = []
-        metrics = {}
-        target = inputs[0]
-        output = inputs[1]
-
-        target_group, target_slot, target_all = raven_utils.decode.decode_target(target)
-
-        group_output, output_slot, outputs = raven_utils.decode.output_divide(output, split_fn=tf.split)
-
-        # group
-        if self.group_loss:
-            group_loss = self.lw * self.plw[0] *  self.loss_fn(target_group, group_output)
-            losses.append(group_loss)
-
-            if isinstance(self.enable_metrics, str):
-                group_metric = self.metric_fn_group(target_group, group_output)
-                # metrics[GROUP] = group_metric
-                self.add_metric(group_metric)
-                self.add_metric(tf.reduce_sum(group_metric), f"{self.enable_metrics}{ACC}")
-
-        # setting uniformity mask
-        full_properties_musks = self.mode(target)
-
-        range_mask = self.range_mask(target_group)
-
-        if self.slot_loss:
-            # number
-            number_mask = range_mask & full_properties_musks[0]
-            number_mask = tf.cast(number_mask, tf.float32)
-            target_number = tf.reduce_sum(
-                tf.cast(target_slot, "float32") * number_mask, axis=-1)
-            output_number = tf.reduce_sum(
-                tf.cast(tf.sigmoid(output_slot) >= 0.5, "float32") * number_mask, axis=-1)
-
-            # output_number = tf.reduce_sum(tf.sigmoid(output_slot) * number_mask, axis=-1)
-            if self.number_loss:
-                scale = 1 / 9
-                if self.number_loss == 2:
-                    output_number_2 = tf.reduce_sum(tf.sigmoid(output_slot) * number_mask, axis=-1)
-                else:
-                    output_number_2 = output_number
-                number_loss = self.lw * self.plw[1] * mse(tf.stop_gradient(target_number) * scale, output_number_2 * scale)
-                losses.append(number_loss)
-
-            # metrics[NUMBER] = number_acc
-
-            if isinstance(self.enable_metrics, str):
-                number_acc = tf.reduce_mean(
-                    tf.cast(tf.cast(target_number, "int8") == tf.cast(output_number, "int8"), "float32"))
-                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}_{NUMBER}")
-                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}")
-                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
-
-            # position/slot
-            slot_mask = range_mask & full_properties_musks[1]
-            # tf.boolean_mask(target_slot,slot_mask)
-
-            if tf.reduce_any(slot_mask):
-                # if tf.reduce_mean(tf.cast(slot_mask, dtype=tf.int32)) > 0:
-                target_slot_masked = tf.boolean_mask(target_slot, slot_mask)[:, None]
-                output_slot_masked = tf.boolean_mask(output_slot, slot_mask)[:, None]
-                loss_slot = self.lw * self.plw[2] * tf.reduce_mean(
-                    self.loss_fn_2(tf.cast(target_slot_masked, "float32"), output_slot_masked))
-                if isinstance(self.enable_metrics, str):
-                    acc_slot = self.metric_fn_2(target_slot_masked, output_slot_masked)
-                    self.add_metric(acc_slot)
-                    self.add_metric(tf.reduce_sum(acc_slot), f"{self.enable_metrics}{ACC}")
-                    self.add_metric(tf.reduce_sum(acc_slot), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
-            else:
-                loss_slot = 0.0
-                acc_slot = -1.0
-
-            losses.append(loss_slot)
-            # metrics[SLOT] = acc_slot
-        # if loss_slot != 0:
-
-        # if tf.reduce_any(slot_mask):
-
-        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}_{NUMBER}")
-        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}")
-        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}_NO_{GROUP}")
-
-        # properties
-        for i, out in enumerate(outputs):
-            shape = (-1, rv.entity.SUM, rv.properties.RAW_SIZE[i])
-            out_reshaped = tf.reshape(out, shape)
-            properties_mask = tf.cast(target_slot, "bool") & full_properties_musks[i + 2]
-
-            if tf.reduce_any(properties_mask):
-                out_masked = tf.boolean_mask(out_reshaped, properties_mask)
-                out_target = tf.boolean_mask(target_all[i], properties_mask)
-                loss = self.lw * self.plw[3+i] * self.loss_fn(out_target, out_masked)
-                if isinstance(self.enable_metrics, str):
-                    metric = self.metric_fn[i](out_target, out_masked)
-                    self.add_metric(metric)
-                    # self.add_metric(metric, f"{self.enable_metrics}{ACC}")
-                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}")
-                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}_{PROPERTIES}")
-                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
-            else:
-                loss = 0.0
-                metric = -1.0
-
-            losses.append(loss)
-        return losses
-
-
-class FullMask(Model):
-    def __init__(self, mode=create_uniform_mask):
-        super().__init__()
-        self.range_mask = RangeMask()
-        self.mode = mode
-
-    def call(self, inputs):
-        target_group, target_slot, _ = raven_utils.decode.decode_target(inputs)
-        full_properties_musks = self.mode(inputs)
-        range_mask = self.range_mask(target_group)
-
-        number_mask = range_mask & full_properties_musks[0]
-
-        slot_mask = range_mask & full_properties_musks[1]
-        properties_mask = []
-        for property_mask in full_properties_musks[2:]:
-            properties_mask.append(tf.cast(target_slot, "bool") & property_mask)
-        return [slot_mask, properties_mask, number_mask]
-
-
-def create_mask(rules, i):
-    mask_1 = tf.tile(rules[:, i][None], [len(rv.target.FIRST_LAYOUT), 1])
-    mask_2 = tf.tile(rules[:, i + 5][None], [len(rv.target.SECOND_LAYOUT), 1])
-    shape = tf.shape(rules)
-    full_mask_1 = tf.scatter_nd(tnp.array(rv.target.FIRST_LAYOUT)[:, None], mask_1, shape=(rv.entity.SUM, shape[0]))
-    full_mask_2 = tf.tensor_scatter_nd_update(full_mask_1, tnp.array(rv.target.SECOND_LAYOUT)[:, None], mask_2)
-    return tf.transpose(full_mask_2)
-
-
-# class PredictModel(Model):
-#     def __init__(self):
-#         super().__init__()
-#         self.predict_fn = Lambda(partial(tf.argmax, axis=-1))
-#         self.predict_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
-#         self.range_mask = RangeMask()
-#
-#     # self.predict_fn = partial(tf.argmax, axis=-1)
-#
-#     def call(self, inputs):
-#         group_output = inputs[rv.OUTPUT_GROUP_SLICE]
-#         group_loss = self.predict_fn(group_output)[:, None]
-#
-#         output_slot = inputs[rv.OUTPUT_SLOT_SLICE]
-#         range_mask = self.range_mask(group_loss[:, 0])
-#         loss_slot = tf.cast(self.predict_fn_2(output_slot), dtype=tf.int64)
-#
-#         properties_output = inputs[rv.OUTPUT_PROPERTIES_SLICE]
-#         properties = []
-#         outputs = tf.split(properties_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-1)
-#         for i, out in enumerate(outputs):
-#             shape = (-1, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
-#             out_reshaped = tf.reshape(out, shape)
-#             properties.append(self.predict_fn(out_reshaped))
-#         number_loss = tf.reduce_sum(loss_slot, axis=-1, keepdims=True)
-#
-#         result = tf.concat([group_loss, loss_slot, interleave(properties), number_loss], axis=-1)
-#
-#         return [result, range_mask, range_mask, range_mask, range_mask]
-
-class PredictModel(Model):
-    def __init__(self):
-        super().__init__()
-        self.predict_fn = Predict()
-        self.predict_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
-        self.range_mask = RangeMask()
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    def call(self, inputs):
-        group_output, output_slot, *properties = rv.decode.output(inputs, tf.split, self.predict_fn, self.predict_fn_2)
-        number_loss = K.int64(K.sum(output_slot))
-        result = tf.concat(
-            [group_output[:, None], tf.cast(output_slot, dtype=tf.int64), interleave(properties), number_loss[:, None]],
-            axis=-1)
-
-        range_mask = self.range_mask(group_output)
-        return [result, range_mask]
-        # return [result, range_mask, range_mask, range_mask, range_mask]
-
-
-# todo change slices
-class PredictModelMasked(Model):
-    def __init__(self):
-        super().__init__()
-        self.predict_fn = Lambda(partial(tf.argmax, axis=-1))
-        self.loss_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
-        self.range_mask = RangeMask()
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    def call(self, inputs):
-        group_output = inputs[:, -rv.GROUPS_NO:]
-        group_loss = self.predict_fn(group_output)[:, None]
-
-        output_slot = inputs[:, :rv.ENTITY_SUM]
-        range_mask = self.range_mask(group_loss[:, 0])
-        loss_slot = tf.cast(self.predict_fn_2(output_slot * range_mask), dtype=tf.int64)
-
-        properties_output = inputs[:, rv.ENTITY_SUM:-rv.GROUPS_NO]
-
-        properties = []
-        outputs = tf.split(properties_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-1)
-        for i, out in enumerate(outputs):
-            shape = (-1, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
-            out_reshaped = tf.reshape(out, shape)
-            out_masked = out_reshaped * loss_slot[..., None]
-            properties.append(self.predict_fn(out_masked))
-            # out_masked[0].numpy()
-        number_loss = tf.reduce_sum(loss_slot, axis=-1, keepdims=True)
-
-        result = tf.concat([group_loss, loss_slot, interleave(properties), number_loss], axis=-1)
-
-        return result
-
-
-def final_predict_mask(x, mask):
-    r = reshape(x[0][:, rv.INDEX[0]:-1], [-1, 3])
-    return tf.ragged.boolean_mask(r, mask)
-
-
-def final_predict(x, mode=False):
-    m = x[1] if mode else tf.cast(x[0][:, 1:rv.INDEX[0]], tf.bool)
-    return final_predict_mask(x[0], m)
-
-
-def final_predict_2(x):
-    ones = tf.cast(tf.ones(tf.shape(x[0])[0]), tf.bool)[:, None]
-    mask = tf.concat([ones, tf.tile(x[1], [1, 4]), ones], axis=-1)
-    return tf.ragged.boolean_mask(x[0], mask)
-
-
-class PredictModelOld(Model):
-
-    def call(self, inputs):
-        output = inputs[-2]
-
-        rest_output = output[:, :-rv.GROUPS_NO]
-
-        result_all = []
-        outputs = tf.split(rest_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-3)
-        for i, out in enumerate(outputs):
-            shape = (-3, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
-            out_reshaped = tf.reshape(out, shape)
-
-            result = tf.cast(tf.argmax(out_reshaped, axis=-3), dtype="int8")
-            result_all.append(result)
-
-        result_all = interleave(result_all)
-        return result_all
-
-
-def get_matches(diff, target_index):
-    diff_sum = K.sum(diff)
-    db_argsort = tf.argsort(diff_sum, axis=-1)
-    db_sorted = tf.sort(diff_sum)
-    db_mask = db_sorted[:, 0, None] == db_sorted
-    db_same = tf.where(db_mask, db_argsort, -1 * tf.ones_like(db_argsort))
-    matched_index = db_same == target_index
-    # setting shape needed for TensorFlow graph
-    matched_index.set_shape(db_same.shape)
-    matches = K.any(matched_index)
-    more_matches = K.sum(db_mask) > 1
-    once_matches = K.sum(matches & tf.math.logical_not(more_matches))
-    return matches, more_matches, once_matches
-
-
-class SimilarityRaven(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False):
-        super().__init__()
-        self.range_mask = RangeMask()
-        self.mode = mode
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    # INDEX, PREDICT, LABELS
-    def call(self, inputs):
-        metrics = []
-        target_index = inputs[0] - 8
-        predict = inputs[1]
-        answers = inputs[2][:, 8:]
-        shape = tf.shape(predict)
-
-        target = K.gather(answers, target_index[:, 0])
-
-        target_group = target[:, 0]
-
-        # comp_slice = np.
-        target_comp = target[:, 1:rv.target.END_INDEX]
-        predict_comp = predict[:, 1:rv.target.END_INDEX]
-        answers_comp = answers[:, :, 1:rv.target.END_INDEX]
-
-        full_properties_musks = self.mode(target)
-        fpm = K.cat([full_properties_musks[0], interleave(full_properties_musks[2:])])
-
-        range_mask = self.range_mask(target_group)
-        full_range_mask = K.cat([range_mask, tf.repeat(range_mask, 3, axis=-1)], axis=-1)
-
-        final_mask = fpm & full_range_mask
-
-        target_masked = target_comp * final_mask
-        predict_masked = predict_comp * final_mask
-        answers_masked = answers_comp * tf.tile(final_mask[:, None], [1, 8, 1])
-
-        acc_same = K.mean(K.all(target_masked == predict_masked))
-        self.add_metric(acc_same, ACC_SAME)
-        metrics.append(acc_same)
-
-        diff = tf.abs(predict_masked[:, None] - answers_masked)
-        diff_bool = diff != 0
-
-        matches, more_matches, once_matches = get_matches(tf.cast(diff_bool, dtype=tf.int32), target_index)
-
-        second_phase_mask = (more_matches & matches)
-        diff_second_phase = tf.boolean_mask(diff, second_phase_mask)
-        target_index_2 = tf.boolean_mask(target_index, second_phase_mask, axis=0)
-
-        matches_2, more_matches_2, once_matches_2 = get_matches(diff_second_phase, target_index_2)
-        matches_2_no = K.sum(matches_2)
-
-        acc_choose_upper = (once_matches + matches_2_no) / shape[0]
-        self.add_metric(acc_choose_upper, ACC_CHOOSE_UPPER)
-        metrics.append(acc_choose_upper)
-
-        acc_choose_lower = (once_matches + once_matches_2) / shape[0]
-        self.add_metric(acc_choose_lower, ACC_CHOOSE_LOWER)
-        metrics.append(acc_choose_lower)
-
-        return metrics
-
-
-class SimilarityRaven2(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False):
-        super().__init__()
-        self.range_mask = RangeMask()
-        self.mode = mode
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    # INDEX, PREDICT, LABELS
-    def call(self, inputs):
-        metrics = []
-        target_index = inputs[0] - 8
-        predict = inputs[1]
-        answers = inputs[2][:, 8:]
-        shape = tf.shape(predict)
-
-        target = K.gather(answers, target_index[:, 0])
-
-        target_group = target[:, 0]
-
-        # comp_slice = np.
-        target_comp = target[:, 1:rv.target.END_INDEX]
-        predict_comp = predict[:, 1:rv.target.END_INDEX]
-        answers_comp = answers[:, :, 1:rv.target.END_INDEX]
-
-        full_properties_musks = self.mode(target)
-        fpm = K.cat([full_properties_musks[0], interleave(full_properties_musks[2:])])
-
-        range_mask = self.range_mask(target_group)
-        full_range_mask = K.cat([range_mask, tf.repeat(range_mask, 3, axis=-1)], axis=-1)
-
-        final_mask = fpm & full_range_mask
-
-        target_masked = target_comp * final_mask
-        predict_masked = predict_comp * final_mask
-        answers_masked = answers_comp * tf.tile(final_mask[:, None], [1, 8, 1])
-
-        acc_same = K.mean(K.all(target_masked == predict_masked))
-        self.add_metric(acc_same, ACC_SAME)
-        metrics.append(acc_same)
-
-        diff = tf.abs(predict_masked[:, None] - answers_masked)
-        diff_bool = diff != 0
-
-        matches, more_matches, once_matches = get_matches(tf.cast(diff_bool, dtype=tf.int32), target_index)
-
-        second_phase_mask = (more_matches & matches)
-        diff_second_phase = tf.boolean_mask(diff, second_phase_mask)
-        target_index_2 = tf.boolean_mask(target_index, second_phase_mask, axis=0)
-
-        matches_2, more_matches_2, once_matches_2 = get_matches(diff_second_phase, target_index_2)
-        matches_2_no = K.sum(matches_2)
-
-        acc_choose_upper = (once_matches + matches_2_no) / shape[0]
-        self.add_metric(acc_choose_upper, ACC_CHOOSE_UPPER)
-        metrics.append(acc_choose_upper)
-
-        acc_choose_lower = (once_matches + once_matches_2) / shape[0]
-        self.add_metric(acc_choose_lower, ACC_CHOOSE_LOWER)
-        metrics.append(acc_choose_lower)
-
-        metrics.append(K.sum(target_masked != predict_masked))
-
-        return metrics
-
-
-class LatentLossModel(Model):
-    def __init__(self, dir_=HORIZONTAL):
-        super().__init__()
-        # self.sum_metrics = []
-        # for i in range(8):
-        #     self.sum_metrics.append(Sum(name=f"no_{i}"))
-        self.metric_fn = Accuracy(name="acc_latent")
-        if dir_ == VERTICAL:
-            self.dir = (6, 7)
-        else:
-            self.dir = (2, 5)
-
-    def call(self, inputs):
-        target_image = tf.reshape(inputs[0][2], [-1])
-        output = inputs[1]
-        latents = tnp.asarray(inputs[2])
-
-        target_hor = tf.concat([
-            latents[:, self.dir],
-            latents[tf.range(latents.shape[0]), target_image + 8][:, None]
-        ],
-            axis=1)
-
-        loss_hor = mse(K.stop_gradient(target_hor), output)
-        self.add_loss(loss_hor)
-
-        self.add_metric(self.metric_fn(inputs[3], target_image))
-
-        return loss_hor
-
-
-class PredRav(Model):
-
-    def call(self, inputs):
-        output = inputs[0][:, -1]
-        answers = inputs[1][:, 8:]
-        return tf.argmin(tf.reduce_sum(tf.abs(output[:, None] - answers), axis=-1), axis=-1)

raven_utils/models/loss_3.py

@@ -1,638 +0,0 @@
-from functools import partial
-
-import tensorflow as tf
-import tensorflow.experimental.numpy as tnp
-from models_utils import OUTPUT, TARGET, PREDICT, DictModel, add_loss, LOSS, Predict
-from models_utils import SubClassingModel
-from models_utils.models.utils import interleave
-from models_utils.op import reshape
-from tensorflow.keras import Model
-# from tensorflow.keras import backend as K
-from tensorflow.keras.layers import Lambda
-from tensorflow.keras.losses import SparseCategoricalCrossentropy, mse
-from tensorflow.keras.metrics import SparseCategoricalAccuracy, Accuracy, BinaryAccuracy
-import models_utils.ops as K
-
-import raven_utils.decode
-import raven_utils as rv
-from raven_utils.config.constant import LABELS, INDEX, ACC_SAME, ACC_CHOOSE_LOWER, ACC_CHOOSE_UPPER, CLASSIFICATION, \
-    SLOT, \
-    PROPERTIES, ACC, GROUP, NUMBER, MASK
-from raven_utils.models.uitls_ import RangeMask
-from raven_utils.const import VERTICAL, HORIZONTAL
-
-
-def get_properties_mask(target):
-    return target[:, rv.target.END_INDEX:rv.target.UNIFORMITY_INDEX] > 0
-
-
-def create_change_mask(target):
-    properties_mask = get_properties_mask(target)
-    return [create_mask(properties_mask, i) for i, _ in enumerate(rv.rules.ATTRIBUTES)]
-
-
-def create_uniform_mask(target):
-    u_mask = lambda i: tf.tile(target[:, rv.target.UNIFORMITY_INDEX + i, None] == 3, [1, rv.rules.ATTRIBUTES_LEN])
-    properties_mask = tf.concat([u_mask(0), u_mask(1)], axis=-1) | get_properties_mask(target)
-    return [create_mask(properties_mask, i) for i, _ in enumerate(rv.rules.ATTRIBUTES)]
-
-
-def create_all_mask(target):
-    return [
-        tf.cast(tf.ones(tf.stack([tf.shape(target)[0], rv.entity.SUM])), dtype=tf.bool) for i, _ in
-        enumerate(rv.rules.ATTRIBUTES)]
-
-
-class BaselineClassificationLossModel(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True):
-        super().__init__()
-        self.predict_fn = SubClassingModel([lambda x: x[0], PredictModel()])
-        self.loss_fn = ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                       group_loss=group_loss)
-        self.metric_fn = SimilarityRaven(mode=mode)
-
-    def call(self, inputs):
-        losses = []
-        output = inputs[1]
-        losses.append(self.loss_fn([inputs[0][0], output]))
-        losses.append(self.metric_fn([inputs[0][2], inputs[3][0], inputs[0][1][:, 8:]]))
-        return losses
-
-
-class RavenLoss(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.3),
-                 classification=False, trans=True, anneal=False):
-        super().__init__()
-        if anneal:
-            self.weight_scheduler
-        self.classification = classification
-        self.trans = trans
-        self.predict_fn = DictModel(SubClassingModel([lambda x: x[-1], PredictModel()]), in_=OUTPUT,
-                                    out=[PREDICT, MASK], name="pred")
-        if self.trans:
-            self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                                    group_loss=group_loss, enable_metrics=False, lw=lw[0]),
-                                    name="main_loss")
-            self.loss_fn_2 = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                                      group_loss=group_loss), name="add_loss")
-            self.metric_fn = SimilarityRaven(mode=mode)
-        if self.classification:
-            self.loss_fn_3 = add_loss(
-                ClassRavenModel(mode=create_all_mask, number_loss=number_loss, slot_loss=slot_loss,
-                                group_loss=group_loss, enable_metrics="c" if self.trans else True), lw=lw[1],
-                name="class_loss")
-
-    def call(self, inputs):
-        losses = []
-        output = inputs[OUTPUT]
-        target = inputs[TARGET]
-        labels = inputs[LABELS]
-
-        if self.trans:
-            losses.append(self.loss_fn([labels[:, 2], output[0]]))
-            losses.append(self.loss_fn([labels[:, 5], output[1]]))
-            losses.append(self.loss_fn_2([target, output[2]]))
-            losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
-        if self.classification:
-            for i in range(8):
-                losses.append(self.loss_fn_3([labels[:, i], inputs[CLASSIFICATION][i]]))
-        return {**inputs, LOSS: losses}
-
-
-class VTRavenLoss(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(2.0, 1.0),
-                 classification=False, trans=True, anneal=False, plw=None):
-        super().__init__()
-        if anneal:
-            self.weight_scheduler
-        self.classification = classification
-        self.trans = trans
-        self.predict_fn = DictModel(SubClassingModel([lambda x: x[:, -1], PredictModel()]), in_=OUTPUT,
-                                    out=[PREDICT, "predict_mask"], name="pred")
-        self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                                group_loss=group_loss, plw=plw), lw=lw[0], name="add_loss")
-        self.metric_fn = SimilarityRaven(mode=mode)
-        if self.classification:
-            self.loss_fn_2 = add_loss(
-                ClassRavenModel(mode=create_all_mask, number_loss=number_loss, slot_loss=slot_loss,
-                                group_loss=group_loss, enable_metrics="c", plw=plw), lw=lw[1], name="class_loss")
-
-    def call(self, inputs):
-        losses = []
-        output = inputs[OUTPUT]
-        target = inputs[TARGET]
-        labels = inputs[LABELS]
-        mask = inputs[MASK]
-
-        target_masked = target[mask]
-        output_masked = output[mask]
-        losses.append(self.loss_fn([target_masked, output_masked]))
-
-        target_unmasked = target[~mask]
-        output_unmasked = output[~mask]
-        losses.append(self.loss_fn_2([target_unmasked, output_unmasked]))
-
-        losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
-        return {**inputs, LOSS: losses}
-
-
-class SingleVTRavenLoss(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.1),
-                 classification=False, trans=True, anneal=False):
-        super().__init__()
-        if anneal:
-            self.weight_scheduler
-        self.classification = classification
-        self.trans = trans
-        self.predict_fn = DictModel(PredictModel(), in_=OUTPUT, out=[PREDICT, MASK], name="pred")
-        self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
-                                                group_loss=group_loss), lw=lw[0], name="add_loss")
-        self.metric_fn = SimilarityRaven(mode=mode)
-
-    def call(self, inputs):
-        losses = []
-        output = inputs[OUTPUT]
-        target = inputs[TARGET]
-        labels = inputs[LABELS]
-
-        losses.append(self.loss_fn([target, output]))
-        losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
-        return {**inputs, LOSS: losses}
-
-
-class ClassRavenModel(Model):
-    def __init__(self, mode=create_all_mask, plw=None, number_loss=False, slot_loss=True, group_loss=True,
-                 enable_metrics=True,
-                 lw=1.0):
-        super().__init__()
-        self.number_loss = number_loss
-        self.group_loss = group_loss
-        self.enable_metrics = enable_metrics
-        self.slot_loss = slot_loss
-        self.predict_fn = PredictModel()
-        self.loss_fn = SparseCategoricalCrossentropy(from_logits=True)
-        if self.slot_loss:
-            self.loss_fn_2 = tf.nn.sigmoid_cross_entropy_with_logits
-        if self.enable_metrics:
-            self.enable_metrics = f"{self.enable_metrics}_" if isinstance(self.enable_metrics, str) else ""
-            self.metric_fn = [
-                SparseCategoricalAccuracy(name=f"{self.enable_metrics}{ACC}_{property_}") for property_ in
-                rv.properties.NAMES]
-            if self.group_loss:
-                self.metric_fn_group = SparseCategoricalAccuracy(name=f"{self.enable_metrics}{ACC}_{GROUP}")
-            if self.slot_loss:
-                self.metric_fn_2 = BinaryAccuracy(name=f"{self.enable_metrics}{ACC}_{SLOT}")
-        self.range_mask = RangeMask()
-        self.mode = mode
-        self.lw = lw
-        if not plw:
-            plw = [1., 95.37352927, 2.83426987, 0.85212836, 1.096005, 1.21943385]
-        elif isinstance(plw, int) or isinstance(plw, float):
-            plw = [1., plw, 2.83426987, 0.85212836, 1.096005, 1.21943385]
-            # plw = [plw] * 6
-        self.plw = plw
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    def call(self, inputs):
-        losses = []
-        metrics = {}
-        target = inputs[0]
-        output = inputs[1]
-
-        target_group, target_slot, target_all = raven_utils.decode.decode_target(target)
-
-        group_output, output_slot, outputs = raven_utils.decode.output_divide(output, split_fn=tf.split)
-
-        # group
-        if self.group_loss:
-            group_loss = self.lw * self.plw[0] * self.loss_fn(target_group, group_output)
-            losses.append(group_loss)
-
-            if isinstance(self.enable_metrics, str):
-                group_metric = self.metric_fn_group(target_group, group_output)
-                # metrics[GROUP] = group_metric
-                self.add_metric(group_metric)
-                self.add_metric(tf.reduce_sum(group_metric), f"{self.enable_metrics}{ACC}")
-
-        # setting uniformity mask
-        full_properties_musks = self.mode(target)
-
-        range_mask = self.range_mask(target_group)
-
-        if self.slot_loss:
-            # number
-            number_mask = range_mask & full_properties_musks[0]
-            number_mask = tf.cast(number_mask, tf.float32)
-            target_number = tf.reduce_sum(
-                tf.cast(target_slot, "float32") * number_mask, axis=-1)
-            output_number = tf.reduce_sum(
-                tf.cast(tf.sigmoid(output_slot) >= 0.5, "float32") * number_mask, axis=-1)
-
-            # output_number = tf.reduce_sum(tf.sigmoid(output_slot) * number_mask, axis=-1)
-            if self.number_loss:
-                scale = 1 / 9
-                if self.number_loss == 2:
-                    output_number_2 = tf.reduce_sum(tf.sigmoid(output_slot) * number_mask, axis=-1)
-                else:
-                    output_number_2 = output_number
-                number_loss = self.lw * self.plw[1] * mse(tf.stop_gradient(target_number) * scale,
-                                                          output_number_2 * scale)
-                losses.append(number_loss)
-
-            # metrics[NUMBER] = number_acc
-
-            if isinstance(self.enable_metrics, str):
-                number_acc = tf.reduce_mean(
-                    tf.cast(tf.cast(target_number, "int8") == tf.cast(output_number, "int8"), "float32"))
-                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}_{NUMBER}")
-                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}")
-                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
-
-            # position/slot
-            slot_mask = range_mask & full_properties_musks[1]
-            # tf.boolean_mask(target_slot,slot_mask)
-
-            if tf.reduce_any(slot_mask):
-                # if tf.reduce_mean(tf.cast(slot_mask, dtype=tf.int32)) > 0:
-                target_slot_masked = tf.boolean_mask(target_slot, slot_mask)[:, None]
-                output_slot_masked = tf.boolean_mask(output_slot, slot_mask)[:, None]
-                loss_slot = self.lw * self.plw[2] * tf.reduce_mean(
-                    self.loss_fn_2(tf.cast(target_slot_masked, "float32"), output_slot_masked))
-                if isinstance(self.enable_metrics, str):
-                    acc_slot = self.metric_fn_2(target_slot_masked, output_slot_masked)
-                    self.add_metric(acc_slot)
-                    self.add_metric(tf.reduce_sum(acc_slot), f"{self.enable_metrics}{ACC}")
-                    self.add_metric(tf.reduce_sum(acc_slot), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
-            else:
-                loss_slot = 0.0
-                acc_slot = -1.0
-
-            losses.append(loss_slot)
-            # metrics[SLOT] = acc_slot
-        # if loss_slot != 0:
-
-        # if tf.reduce_any(slot_mask):
-
-        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}_{NUMBER}")
-        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}")
-        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}_NO_{GROUP}")
-
-        # properties
-        for i, out in enumerate(outputs):
-            shape = (-1, rv.entity.SUM, rv.properties.RAW_SIZE[i])
-            out_reshaped = tf.reshape(out, shape)
-            properties_mask = tf.cast(target_slot, "bool") & full_properties_musks[i + 2]
-
-            if tf.reduce_any(properties_mask):
-                out_masked = tf.boolean_mask(out_reshaped, properties_mask)
-                out_target = tf.boolean_mask(target_all[i], properties_mask)
-                loss = self.lw * self.plw[3 + i] * self.loss_fn(out_target, out_masked)
-                if isinstance(self.enable_metrics, str):
-                    metric = self.metric_fn[i](out_target, out_masked)
-                    self.add_metric(metric)
-                    # self.add_metric(metric, f"{self.enable_metrics}{ACC}")
-                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}")
-                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}_{PROPERTIES}")
-                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
-            else:
-                loss = 0.0
-                metric = -1.0
-
-            losses.append(loss)
-        return losses
-
-
-class FullMask(Model):
-    def __init__(self, mode=create_uniform_mask):
-        super().__init__()
-        self.range_mask = RangeMask()
-        self.mode = mode
-
-    def call(self, inputs):
-        target_group, target_slot, _ = raven_utils.decode.decode_target(inputs)
-        full_properties_musks = self.mode(inputs)
-        range_mask = self.range_mask(target_group)
-
-        number_mask = range_mask & full_properties_musks[0]
-
-        slot_mask = range_mask & full_properties_musks[1]
-        properties_mask = []
-        for property_mask in full_properties_musks[2:]:
-            properties_mask.append(tf.cast(target_slot, "bool") & property_mask)
-        return [slot_mask, properties_mask, number_mask]
-
-
-def create_mask(rules, i):
-    mask_1 = tf.tile(rules[:, i][None], [len(rv.target.FIRST_LAYOUT), 1])
-    mask_2 = tf.tile(rules[:, i + 5][None], [len(rv.target.SECOND_LAYOUT), 1])
-    shape = tf.shape(rules)
-    full_mask_1 = tf.scatter_nd(tnp.array(rv.target.FIRST_LAYOUT)[:, None], mask_1, shape=(rv.entity.SUM, shape[0]))
-    full_mask_2 = tf.tensor_scatter_nd_update(full_mask_1, tnp.array(rv.target.SECOND_LAYOUT)[:, None], mask_2)
-    return tf.transpose(full_mask_2)
-
-
-# class PredictModel(Model):
-#     def __init__(self):
-#         super().__init__()
-#         self.predict_fn = Lambda(partial(tf.argmax, axis=-1))
-#         self.predict_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
-#         self.range_mask = RangeMask()
-#
-#     # self.predict_fn = partial(tf.argmax, axis=-1)
-#
-#     def call(self, inputs):
-#         group_output = inputs[rv.OUTPUT_GROUP_SLICE]
-#         group_loss = self.predict_fn(group_output)[:, None]
-#
-#         output_slot = inputs[rv.OUTPUT_SLOT_SLICE]
-#         range_mask = self.range_mask(group_loss[:, 0])
-#         loss_slot = tf.cast(self.predict_fn_2(output_slot), dtype=tf.int64)
-#
-#         properties_output = inputs[rv.OUTPUT_PROPERTIES_SLICE]
-#         properties = []
-#         outputs = tf.split(properties_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-1)
-#         for i, out in enumerate(outputs):
-#             shape = (-1, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
-#             out_reshaped = tf.reshape(out, shape)
-#             properties.append(self.predict_fn(out_reshaped))
-#         number_loss = tf.reduce_sum(loss_slot, axis=-1, keepdims=True)
-#
-#         result = tf.concat([group_loss, loss_slot, interleave(properties), number_loss], axis=-1)
-#
-#         return [result, range_mask, range_mask, range_mask, range_mask]
-
-class PredictModel(Model):
-    def __init__(self):
-        super().__init__()
-        self.predict_fn = Predict()
-        self.predict_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
-        self.range_mask = RangeMask()
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    def call(self, inputs):
-        group_output, output_slot, *properties = rv.decode.output(inputs, tf.split, self.predict_fn, self.predict_fn_2)
-        number_loss = K.int64(K.sum(output_slot))
-        result = tf.concat(
-            [group_output[:, None], tf.cast(output_slot, dtype=tf.int64), interleave(properties), number_loss[:, None]],
-            axis=-1)
-
-        range_mask = self.range_mask(group_output)
-        return [result, range_mask]
-        # return [result, range_mask, range_mask, range_mask, range_mask]
-
-
-# todo change slices
-class PredictModelMasked(Model):
-    def __init__(self):
-        super().__init__()
-        self.predict_fn = Lambda(partial(tf.argmax, axis=-1))
-        self.loss_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
-        self.range_mask = RangeMask()
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    def call(self, inputs):
-        group_output = inputs[:, -rv.GROUPS_NO:]
-        group_loss = self.predict_fn(group_output)[:, None]
-
-        output_slot = inputs[:, :rv.ENTITY_SUM]
-        range_mask = self.range_mask(group_loss[:, 0])
-        loss_slot = tf.cast(self.predict_fn_2(output_slot * range_mask), dtype=tf.int64)
-
-        properties_output = inputs[:, rv.ENTITY_SUM:-rv.GROUPS_NO]
-
-        properties = []
-        outputs = tf.split(properties_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-1)
-        for i, out in enumerate(outputs):
-            shape = (-1, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
-            out_reshaped = tf.reshape(out, shape)
-            out_masked = out_reshaped * loss_slot[..., None]
-            properties.append(self.predict_fn(out_masked))
-            # out_masked[0].numpy()
-        number_loss = tf.reduce_sum(loss_slot, axis=-1, keepdims=True)
-
-        result = tf.concat([group_loss, loss_slot, interleave(properties), number_loss], axis=-1)
-
-        return result
-
-
-def final_predict_mask(x, mask):
-    r = reshape(x[0][:, rv.INDEX[0]:-1], [-1, 3])
-    return tf.ragged.boolean_mask(r, mask)
-
-
-def final_predict(x, mode=False):
-    m = x[1] if mode else tf.cast(x[0][:, 1:rv.INDEX[0]], tf.bool)
-    return final_predict_mask(x[0], m)
-
-
-def final_predict_2(x):
-    ones = tf.cast(tf.ones(tf.shape(x[0])[0]), tf.bool)[:, None]
-    mask = tf.concat([ones, tf.tile(x[1], [1, 4]), ones], axis=-1)
-    return tf.ragged.boolean_mask(x[0], mask)
-
-
-class PredictModelOld(Model):
-
-    def call(self, inputs):
-        output = inputs[-2]
-
-        rest_output = output[:, :-rv.GROUPS_NO]
-
-        result_all = []
-        outputs = tf.split(rest_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-3)
-        for i, out in enumerate(outputs):
-            shape = (-3, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
-            out_reshaped = tf.reshape(out, shape)
-
-            result = tf.cast(tf.argmax(out_reshaped, axis=-3), dtype="int8")
-            result_all.append(result)
-
-        result_all = interleave(result_all)
-        return result_all
-
-
-def get_matches(diff, target_index):
-    diff_sum = K.sum(diff)
-    db_argsort = tf.argsort(diff_sum, axis=-1)
-    db_sorted = tf.sort(diff_sum)
-    db_mask = db_sorted[:, 0, None] == db_sorted
-    db_same = tf.where(db_mask, db_argsort, -1 * tf.ones_like(db_argsort))
-    matched_index = db_same == target_index
-    # setting shape needed for TensorFlow graph
-    matched_index.set_shape(db_same.shape)
-    matches = K.any(matched_index)
-    more_matches = K.sum(db_mask) > 1
-    once_matches = K.sum(matches & tf.math.logical_not(more_matches))
-    return matches, more_matches, once_matches
-
-
-class SimilarityRaven(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False):
-        super().__init__()
-        self.range_mask = RangeMask()
-        self.mode = mode
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    # INDEX, PREDICT, LABELS
-    def call(self, inputs):
-        metrics = []
-        target_index = inputs[0] - 8
-        predict = inputs[1]
-        answers = inputs[2][:, 8:]
-        shape = tf.shape(predict)
-
-        target = K.gather(answers, target_index[:, 0])
-
-        target_group = target[:, 0]
-
-        # comp_slice = np.
-        target_comp = target[:, 1:rv.target.END_INDEX]
-        predict_comp = predict[:, 1:rv.target.END_INDEX]
-        answers_comp = answers[:, :, 1:rv.target.END_INDEX]
-
-        full_properties_musks = self.mode(target)
-        fpm = K.cat([full_properties_musks[0], interleave(full_properties_musks[2:])])
-
-        range_mask = self.range_mask(target_group)
-        full_range_mask = K.cat([range_mask, tf.repeat(range_mask, 3, axis=-1)], axis=-1)
-
-        final_mask = fpm & full_range_mask
-
-        target_masked = target_comp * final_mask
-        predict_masked = predict_comp * final_mask
-        answers_masked = answers_comp * tf.tile(final_mask[:, None], [1, 8, 1])
-
-        acc_same = K.mean(K.all(target_masked == predict_masked))
-        self.add_metric(acc_same, ACC_SAME)
-        metrics.append(acc_same)
-
-        diff = tf.abs(predict_masked[:, None] - answers_masked)
-        diff_bool = diff != 0
-
-        matches, more_matches, once_matches = get_matches(tf.cast(diff_bool, dtype=tf.int32), target_index)
-
-        second_phase_mask = (more_matches & matches)
-        diff_second_phase = tf.boolean_mask(diff, second_phase_mask)
-        target_index_2 = tf.boolean_mask(target_index, second_phase_mask, axis=0)
-
-        matches_2, more_matches_2, once_matches_2 = get_matches(diff_second_phase, target_index_2)
-        matches_2_no = K.sum(matches_2)
-
-        acc_choose_upper = (once_matches + matches_2_no) / shape[0]
-        self.add_metric(acc_choose_upper, ACC_CHOOSE_UPPER)
-        metrics.append(acc_choose_upper)
-
-        acc_choose_lower = (once_matches + once_matches_2) / shape[0]
-        self.add_metric(acc_choose_lower, ACC_CHOOSE_LOWER)
-        metrics.append(acc_choose_lower)
-
-        return metrics
-
-
-class SimilarityRaven2(Model):
-    def __init__(self, mode=create_all_mask, number_loss=False):
-        super().__init__()
-        self.range_mask = RangeMask()
-        self.mode = mode
-
-    # self.predict_fn = partial(tf.argmax, axis=-1)
-
-    # INDEX, PREDICT, LABELS
-    def call(self, inputs):
-        metrics = []
-        target_index = inputs[0] - 8
-        predict = inputs[1]
-        answers = inputs[2][:, 8:]
-        shape = tf.shape(predict)
-
-        target = K.gather(answers, target_index[:, 0])
-
-        target_group = target[:, 0]
-
-        # comp_slice = np.
-        target_comp = target[:, 1:rv.target.END_INDEX]
-        predict_comp = predict[:, 1:rv.target.END_INDEX]
-        answers_comp = answers[:, :, 1:rv.target.END_INDEX]
-
-        full_properties_musks = self.mode(target)
-        fpm = K.cat([full_properties_musks[0], interleave(full_properties_musks[2:])])
-
-        range_mask = self.range_mask(target_group)
-        full_range_mask = K.cat([range_mask, tf.repeat(range_mask, 3, axis=-1)], axis=-1)
-
-        final_mask = fpm & full_range_mask
-
-        target_masked = target_comp * final_mask
-        predict_masked = predict_comp * final_mask
-        answers_masked = answers_comp * tf.tile(final_mask[:, None], [1, 8, 1])
-
-        acc_same = K.mean(K.all(target_masked == predict_masked))
-        self.add_metric(acc_same, ACC_SAME)
-        metrics.append(acc_same)
-
-        diff = tf.abs(predict_masked[:, None] - answers_masked)
-        diff_bool = diff != 0
-
-        matches, more_matches, once_matches = get_matches(tf.cast(diff_bool, dtype=tf.int32), target_index)
-
-        second_phase_mask = (more_matches & matches)
-        diff_second_phase = tf.boolean_mask(diff, second_phase_mask)
-        target_index_2 = tf.boolean_mask(target_index, second_phase_mask, axis=0)
-
-        matches_2, more_matches_2, once_matches_2 = get_matches(diff_second_phase, target_index_2)
-        matches_2_no = K.sum(matches_2)
-
-        acc_choose_upper = (once_matches + matches_2_no) / shape[0]
-        self.add_metric(acc_choose_upper, ACC_CHOOSE_UPPER)
-        metrics.append(acc_choose_upper)
-
-        acc_choose_lower = (once_matches + once_matches_2) / shape[0]
-        self.add_metric(acc_choose_lower, ACC_CHOOSE_LOWER)
-        metrics.append(acc_choose_lower)
-
-        metrics.append(K.sum(target_masked != predict_masked))
-
-        return metrics
-
-
-class LatentLossModel(Model):
-    def __init__(self, dir_=HORIZONTAL):
-        super().__init__()
-        # self.sum_metrics = []
-        # for i in range(8):
-        #     self.sum_metrics.append(Sum(name=f"no_{i}"))
-        self.metric_fn = Accuracy(name="acc_latent")
-        if dir_ == VERTICAL:
-            self.dir = (6, 7)
-        else:
-            self.dir = (2, 5)
-
-    def call(self, inputs):
-        target_image = tf.reshape(inputs[0][2], [-1])
-        output = inputs[1]
-        latents = tnp.asarray(inputs[2])
-
-        target_hor = tf.concat([
-            latents[:, self.dir],
-            latents[tf.range(latents.shape[0]), target_image + 8][:, None]
-        ],
-            axis=1)
-
-        loss_hor = mse(K.stop_gradient(target_hor), output)
-        self.add_loss(loss_hor)
-
-        self.add_metric(self.metric_fn(inputs[3], target_image))
-
-        return loss_hor
-
-
-class PredRav(Model):
-
-    def call(self, inputs):
-        output = inputs[0][:, -1]
-        answers = inputs[1][:, 8:]
-        return tf.argmin(tf.reduce_sum(tf.abs(output[:, None] - answers), axis=-1), axis=-1)

raven_utils/models/multi_transformer.py

@@ -1,274 +0,0 @@
-import tensorflow as tf
-from functools import partial
-from tensorflow.keras.layers import Lambda
-from tensorflow.keras.layers import Dense
-from tensorflow.keras import Input, Model
-from tensorflow.python.keras import Sequential
-
-from config.constant import TRANS
-from ml_utils import filter_init
-from models.loss import VTRavenLoss, create_uniform_mask, SingleVTRavenLoss
-from models_utils import pmodel, DictModel, bt, INPUTS, bm, OUTPUT, LATENTS, transformer, BatchModel, get_extractor, \
-    build_seq_model, BUILD, build_train_list, InitialWeight
-from models_utils import SumPositionEmbedding, TransformerBlock, CatPositionEmbedding, transformer, BatchInitialWeight
-import models_utils.ops as K
-from models_utils.image import inverse_fn
-from models_utils.ops_core import IndexReshape
-from models_utils.random_ import EpsilonGreedy, EpsilonSoft
-from models_utils.step import StepDict
-
-
-def init_weights(shape, dtype=None):
-    return tf.cast(K.var.image(shape=shape, pre=True), dtype=tf.float32)
-
-
-def conversion(x, max_=45):
-    shape = tf.shape(x)
-    return tf.reshape(x[:, :max_], tf.stack([shape[0], 9, -1]))
-
-
-def take_left(x):
-    return x[..., 7:8]
-
-
-def take_by_index(x, i=8):
-    return x[..., i:i + 1]
-
-
-def mix(x):
-    return (x[..., 7:8] + x[..., 5:6]) / 2
-
-
-def empty_last(x):
-    return tf.zeros_like(x[..., 7:8])
-
-
-class Conversion(Model):
-    def __init__(self):
-        super().__init__()
-        self.model = IndexReshape((0, "9", None))
-
-    def call(self, inputs):
-        return self.model(inputs[:, :45])
-
-
-class RandomImageMask(Model):
-    def __init__(self, last, last_index=9):
-        super().__init__()
-        self.get_last = last
-        self.last_index = last_index
-
-    def call(self, inputs):
-        shape = tf.shape(inputs)
-        indexes = tf.random.uniform(shape=shape[0:1], maxval=self.last_index, dtype=tf.int32)
-        mask = tf.one_hot(indexes, self.last_index)[:, None, None]
-
-        return (1 - mask) * inputs[..., :self.last_index] + mask * tf.tile(self.get_last(inputs),
-                                                                           (1, 1, 1, self.last_index))
-
-
-# res = (1 - mask) * inputs[..., :self.last_index] + mask * tf.tile(self.get_last(inputs),
-#                                                                            (1, 1, 1, self.last_index))
-
-
-# from data_utils import ims
-# for i in range(50):
-#     ims(res[i].numpy().swapaxes(0, 2))
-# res[12].numpy()
-# self.get_last(inputs).numpy()
-# import tensorflow as tf
-# tf.random.uniform(shape=shape[0:1], maxval=255, dtype=tf.int32)
-# from ml_utils import print_error
-# ims(mask[0].numpy())
-# print_error(lambda :ims(mask[0]))
-# from models_utils import ops as K
-
-
-class ImageMask(Model):
-    def __init__(self, last, index=8, last_index=9):
-        super().__init__()
-        self.get_last = last
-        self.index = index
-        self.last_index = last_index
-
-    def call(self, inputs):
-        return tf.concat([inputs[..., :8], self.get_last(inputs)], axis=-1)
-
-
-class CreateGrid(Model):
-    def __init__(self,
-                 no=4,
-                 extractor="ef",
-                 type_=3,
-                 base="seq",
-                 last=take_left,
-                 epsilon=None,
-                 pooling=None,
-                 mask_fn=None,
-                 model=None,
-                 **kwargs
-                 ):
-        super().__init__()
-        self.type_ = type_
-        if type_ == 9:
-            self.start_shape = 75
-            data = (224, 224, 3)
-            conv = lambda: Conversion()
-        else:
-            self.start_shape = 84
-            data = (84, 84, 3)
-            extractor = BUILD[base]([
-                BatchModel(get_extractor(data=data, model=extractor)),
-                lambda x: tf.transpose(x, (1, 0, 2, 3, 4))
-                # lambda x: tf.tile(x[:, :224, :224], (1, 1, 1, 3))
-            ])
-            conv = lambda: conversion
-
-        self.epsilon = epsilon
-        if mask_fn == "random":
-            mask_fn = RandomImageMask(last=last)
-        elif mask_fn is None:
-            mask_fn = ImageMask(last=last)
-
-        self.mask_fn = mask_fn
-
-
-    def call(self, inputs):
-        transposed = tf.image.resize(tf.transpose(inputs, (0, 2, 3, 1)), (self.start_shape, self.start_shape))
-        re = self.mask_fn(transposed)
-
-        # re = tf.concat([transposed[..., :8], self.get_last(transposed)], axis=-1)
-        if self.type_ == 9:
-            x = tf.transpose(re, [0, 3, 1, 2])[..., None]
-            x = K.create_image_grid(x, 3, 3)
-            x = x[:, :224, :224]
-            x = tf.tile(x, [1, 1, 1, 3])
-        else:
-
-            x = tf.stack([
-                re[..., :3],
-                re[..., 3:6],
-                re[..., 6:9],
-            ])
-        return self.model(x)
-
-
-# self.model.layers[0](x)
-
-
-def grid_transformer(
-        *args,
-        type_=9,
-        no=4,
-        extractor="ef",
-        loss_mode=create_uniform_mask,
-        output_size=10,
-        loss_weight=1.0,
-        out_layers=(1000, 1000, 1000),
-        pos_emd="cat",
-        base="seq",
-        inverse_image=True,
-        last="left",
-        mask_fn=None,
-        model=None,
-        trans=None,
-        **kwargs):
-
-    if last == "left":
-        last = take_left
-    elif last == "mix":
-        last = mix
-    elif last == "empty":
-        last = empty_last
-    elif last == "start":
-        last = Sequential([Lambda(empty_last), BatchInitialWeight(initializer=init_weights)])
-
-    create_grid = CreateGrid(
-        type_=type_,
-        no=no,
-        extractor=extractor,
-        model=model,
-        output_size=output_size,
-        out_layer=out_layers,
-        pos_emd=pos_emd,
-        base=base,
-        last=last,
-        mask_fn=mask_fn,
-        **kwargs
-    )
-
-    if model is None:
-        trans = transformer(
-            extractor=extractor,
-            pos_emd=pos_emd,
-            data=data,
-            output_size=output_size,
-            out_layers=out_layer,
-            pooling=conv,
-            no=no,
-            base=base,
-            **kwargs
-            # **as_dict(p.trans)
-        )
-    else:
-        trans = trans
-
-
-
-def get_rav_trans(
-        *args,
-        type_=9,
-        no=4,
-        extractor="ef",
-        loss_mode=create_uniform_mask,
-        output_size=10,
-        loss_weight=1.0,
-        out_layers=(1000, 1000, 1000),
-        pos_emd="cat",
-        base="seq",
-        inverse_image=True,
-        last="left",
-        epsilon="greedy",
-        epsilon_step=500,
-        mask_fn=None,
-        model=None,
-        loss="multi",
-        **kwargs):
-    if last == "left":
-        last = take_left
-    elif last == "mix":
-        last = mix
-    elif last == "empty":
-        last = empty_last
-    elif last == "start":
-        last = Sequential([Lambda(empty_last), BatchInitialWeight(initializer=init_weights)])
-
-    trans_raven = CreateGrid(
-        type_=type_,
-        no=no,
-        extractor=extractor,
-        model=model,
-        output_size=output_size,
-        out_layer=out_layers,
-        pos_emd=pos_emd,
-        base=base,
-        last=last,
-        epsilon=epsilon,
-        mask_fn=mask_fn,
-        **kwargs
-    )
-
-    if loss == "single":
-        loss = SingleVTRavenLoss
-    else:
-        loss = VTRavenLoss
-
-    return bt(
-        DictModel(
-            Sequential([Lambda(lambda x: 255 - x), trans_raven]) if inverse_image else trans_raven,
-            in_=INPUTS,
-            name="Body"
-        ),
-        loss=loss(mode=loss_mode, classification=True, lw=(loss_weight, 1.0)),
-        loss_wrap=False
-    )

raven_utils/models/raven.py

@@ -1,239 +0,0 @@
-from ml_utils import lw, lu
-from models_utils import bm, Base, res, bt, DictModel, dense_drop, drop, build_encoder, MODEL_ARCH, ListModel, short, \
-    dense, Flatten, Cat, CatDenseBefore, \
-    CatDense, CatBefore, Drop, Flat2, down, Pass, conv, Flat, Get, bs, Res, SoftBlock
-from models_utils import SubClassingModel
-from models_utils.config.constants import *
-import numpy as np
-
-from config.constant import *
-from tensorflow.keras.layers import Dense, Activation, BatchNormalization
-import tensorflow as tf
-
-import raven_utils as rv
-
-from models.body import create_block
-from models.class_ import Merge, RavenClass
-from models.head import LatentHeadModel
-
-from models.loss import RavenLoss
-from models.trans import TransModel, FullTrans
-from raven_utils.const import HORIZONTAL
-
-
-def raven_model(scales,
-                out_layers,
-                latent=(64, 128, 256),
-                output_size=None,
-                padding=SAME,
-                body_layers=1,
-                encoder=None,
-                loop=1,
-                model=None,
-                act=None,
-                simpler=0,
-                loss_mode=None,
-                loss_weight=0.3,
-                dir_=HORIZONTAL,
-                global_context=False,
-                images_no=8,
-                context_mul=2,
-                res_act="pass",
-                drop_latent=0,
-                drop_inference=0,
-                drop_end=0,
-                ga=False,
-                trans_norm=None,
-                trans_act="relu",
-                arch=HEAD3,
-                encoder_norm=False,
-                encoder_pool=False,
-                encoder_global="GM",
-                encoder_before=False,
-                tail_units=256,
-                tail_flatten=None,
-                # for now by default
-                tail_down="MP",
-                trans_no=1,
-                trans_score_activation=tf.nn.softmax,
-                block_=SoftBlock,
-                **kwargs):
-    if isinstance(latent, int):
-        latent = (latent, 128, 256)
-    scales = lw(scales)
-
-    context_size = np.array(latent) * context_mul
-    # context_size = latent[scales] * context_mul
-
-    # if scales == 2:
-    #     arch = HEAD
-    # elif scales == 1:
-    #     arch = HEAD2
-    # else:
-    #     arch = VERY2
-
-    if encoder_pool:
-        strides = (1, 1)
-    else:
-        strides = (2, 2)
-    if not isinstance(encoder_before, tuple):
-        encoder_before = [encoder_before] * 3
-
-    # if trans == 1:
-    #     trans_model = TransModel2
-    # else:
-    #     trans_model = TransModel
-
-    # if scales == 3:
-    #     head = MultiHeadModel(encoder=encoder)
-    arch = MODEL_ARCH[arch]
-    heads = []
-    for s in list(range(0, max(scales) + 1)):
-        if s in (0, 1):
-            if s == 0:
-                encoder = build_encoder(arch[:3], add_norm=encoder_norm, add_pool=encoder_pool, kerner_size=(4, 4),
-                                        strides=strides)
-            else:
-                encoder = build_encoder(arch[3:4], add_norm=encoder_norm, add_pool=encoder_pool, kerner_size=(4, 4),
-                                        strides=strides)
-            head = LatentHeadModel(
-                encoder=encoder,
-                inference_network=(
-                    bm([
-                           CatBefore(filters=int(context_size[s] / 8)) if encoder_before[s] else Cat(
-                               filters=context_size[s]),
-                           # todo activation?
-                           Res(filters=context_size[s], padding=padding)
-                       ] + ([drop(drop_inference)] if drop_inference else []),
-                       name="inference")
-                ) if s in scales else Pass(),
-                stem=Base(
-                    bm(
-                        # ok we choose by parameters anyway
-                        [res(filters=latent[s], padding=padding, act=act)] + (
-                            [drop(drop_latent)] if drop_latent else [])
-                    ),
-                    name="stem")
-            )
-        else:
-            encoder = bm([
-                Res(),
-                build_encoder(arch[4:], add_norm=encoder_norm, add_pool=encoder_pool, kerner_size=(4, 4),
-                              strides=strides),
-                short(encoder_global) if encoder_global else Flatten(),
-                dense(latent[s])
-            ])
-            head = LatentHeadModel(
-                encoder=encoder,
-                inference_network=bm([
-                                         # todo Echeck Cat
-                                         CatDenseBefore(filters=int(context_size[s] / 8)) if encoder_before[
-                                             s] else CatDense(filters=context_size[s]),
-                                         # todo activation?
-                                         Res(model="dv2", filters=context_size[s], padding=padding)
-                                     ] + ([dense_drop(drop_inference)] if drop_inference else []),
-                                     name="inference"),
-                stem=Base(
-                    bm(
-                        # ok we choose by parameters anyway
-                        [res(model="dv2", units=latent[s], padding=padding, act=act)] + (
-                            [dense_drop(drop_latent)] if drop_latent else [])
-                    ),
-                    name="stem")
-            )
-        heads.append(head)
-
-    concat_input = [f"{LATENT}_{i}" for i, _ in enumerate(heads)] + [f"{INFERENCE}_{i}" for i, _ in enumerate(heads)]
-    concat_output = ["LATENTS", "INFERENCES"]
-
-    def head_concat(inputs):
-        latents = inputs[:len(heads)]
-        inferences = inputs[len(heads):]
-        return latents, inferences
-
-    head = ListModel([(h, (INPUTS if i == 0 else OUTPUT), [f"{LATENT}_{i}", f"{INFERENCE}_{i}", OUTPUT]) for i, h in
-                      enumerate(heads)] + [
-                         (head_concat, concat_input, concat_output)], out=concat_output)
-    # from rav_utils.raven import init_image
-    # a = init_image()
-    # head(a)
-
-    if model is None:
-        model = []
-        for i in scales:
-            trans_models = []
-            for t in range(trans_no):
-                trans_models.append(
-                    bm(
-                        [create_block(latent=latent[i], simpler=simpler, padding=padding, norm=trans_norm, act=res_act,
-                                      loop=loop, type_="dense" if i == 2 else "conv", block_=block_)] +
-                        [Activation(trans_act)] + [
-                            res(filters=latent[i],
-                                padding=padding,
-                                act=act,
-                                name="body_out",
-                                model="dv2" if i == 2 else "v2") for _ in
-                            range(body_layers)] + ([Drop(drop_latent)] if drop_latent else []),
-                        base_class=SubClassingModel)
-                )
-            trans_models = lu(trans_models)
-            if trans_no > 1:
-                trans_models = bm([
-                    lambda x: [[x[0], x[1]], x[1]],
-                    SoftBlock(
-                        model=trans_models,
-                        score_model=bm([
-                            Flat2(filters=latent[i], units=256, res_no=2),
-                            Dense(trans_no, trans_score_activation)
-                        ])
-                    )
-                ],
-                    base_class=SubClassingModel
-                )
-
-            model.append(
-                TransModel(
-                    body=trans_models,
-                    dir_=dir_,
-                    images_no=images_no
-                )
-            )
-
-    tail = []
-    for i, s in enumerate(scales):
-        flatting = lambda: Flat2(filters=latent[s + 1], base_class=tail_flatten, units=tail_units)
-        if s == 0:
-            if tail_flatten is None:
-                branch = bm([res(filters=latent[s], padding=padding),
-                             conv(filters=latent[s], padding=padding),
-                             BatchNormalization(),
-                             conv(filters=latent[s], padding=padding),
-                             Flatten()])
-            else:
-                branch = bm([down(base_class=tail_down), flatting()])
-        elif s == 1:
-            if tail_flatten is None:
-                branch = bm([res(filters=latent[s], padding=padding),
-                             Flatten()])
-            else:
-                branch = flatting()
-        else:
-            branch = bm([tail_units] * 2, add_flatten=False)
-        tail.append(branch)
-
-    tail.append(
-        bm([dense(tail_units)] + ([dense_drop(drop_end)] if drop_end else []) + [Dense(output_size)], add_flatten=False,
-           name=TAIL))
-    class_input = []
-
-    return bt([
-        DictModel(head, in_=INPUTS, out=[LATENT, INFERENCE], name="Head"),
-        DictModel(FullTrans(model, scales=scales), in_=[LATENT, INFERENCE], out=TRANS, name="Body"),
-        DictModel(RavenClass(Merge(tail), scales=scales, no=8), in_=[LATENT] + class_input, out=CLASSIFICATION,
-                  name="Classificator"),
-        DictModel(RavenClass(Merge(tail), scales=list(range(len(scales))), no=3), in_=[TRANS] + class_input,
-                  out=OUTPUT, name="Classificator_trans"),
-    ],
-        loss=RavenLoss(mode=loss_mode, classification=True, trans=True, lw=(1.0, loss_weight)),
-        loss_wrap=False
-    )

raven_utils/models/trans.py

@@ -1,74 +0,0 @@
-import tensorflow as tf
-from ml_utils import lw
-from models_utils import ops as K, SubClassingModel
-from tensorflow.keras import Model
-
-from models.body import create_dense_block
-import raven_utils as rv
-from raven_utils.const import HORIZONTAL, VERTICAL
-
-
-class TransModel(Model):
-    def __init__(self, body=None, dir_=HORIZONTAL, images_no=8, latent=64):
-        super().__init__()
-        self.model = body or create_dense_block(latent=latent)
-        if dir_ == VERTICAL:
-            self.dir = (0, 3, 1, 4, 3, 5)
-        else:
-            self.dir = (0, 1, 3, 4, 6, 7)
-        self.images_no = images_no
-        self.latent = latent
-
-    def call(self, inputs):
-        # latents = tnp.asarray(inputs[0])
-        latents = inputs[0]
-        inference = inputs[1]
-        shape = tf.shape(latents)
-        new_shape = K.cat([[-1, 3, 2], shape[2:]])
-        horizontal = latents[:, self.dir].reshape(new_shape)
-        res = tf.TensorArray(tf.float32, size=3)
-        for i in range(3):
-            res = res.write(i, self.model([horizontal[:, i], inference]))
-        result = K.tran(res.stack())
-        return result
-
-
-class TransModel2(Model):
-    def __init__(self, body=None, dir_=HORIZONTAL, images_no=8, latent=64):
-        super().__init__()
-        self.body = body or create_dense_block(latent=latent)
-        if dir_ == VERTICAL:
-            self.dir = (0, 3, 1, 4, 3, 5)
-        else:
-            self.dir = (0, 1, 3, 4, 6, 7)
-        self.images_no = images_no
-        self.latent = latent
-
-    def call(self, inputs):
-        # latents = tnp.asarray(inputs[0])
-        latents = inputs[0]
-        inference = inputs[1]
-        shape = tf.shape(latents)
-        new_shape = K.cat([[-1, 3, 2], shape[2:]])
-        horizontal = latents[:, self.dir].reshape(new_shape)
-        res = tf.TensorArray(tf.float32, size=3)
-        for i in tf.range(3):
-            res = res.write(i, self.body([horizontal[:, i], inference[:,i]]))
-        result = K.tran(res.stack())
-        return result
-
-
-class FullTrans(SubClassingModel):
-    def __init__(self, model,scales,name=None):
-        super().__init__(model=model,name=name)
-        self.scales = scales
-
-    def call(self, inputs):
-        latent = lw(inputs[0])
-        inference = lw(inputs[1])
-        results = []
-        # todo merging inference?
-        for i,s in enumerate(self.scales):
-            # results.append(model([latent[::-1][i], inference]))
-            results.append(self.model[i]([latent[s], inference[s]]))
-        return results,

raven_utils/models/transformer.py

@@ -1,133 +0,0 @@
-import tensorflow as tf
-from tensorflow.keras.layers import Lambda
-from tensorflow.python.keras import Sequential
-
-# from models_utils.models.loss import VTRavenLoss, create_uniform_mask, SingleVTRavenLoss
-from models_utils import DictModel, bt, INPUTS, BatchInitialWeight
-import models_utils.ops as K
-from models_utils.models.transformer.img_seq import init_weights, take_left, mix, empty_last
-from models_utils.models.transformer.img_seq2 import init_weights, take_left, mix, empty_last, img_sec_trans
-from models_utils.ops_core import IndexReshape
-from models_utils.random_ import EpsilonGreedy, EpsilonSoft
-from models_utils.step import StepDict
-
-# res = (1 - mask) * inputs[..., :self.last_index] + mask * tf.tile(self.get_last(inputs),
-#                                                                            (1, 1, 1, self.last_index))
-
-
-# from data_utils import ims
-# for i in range(50):
-#     ims(res[i].numpy().swapaxes(0, 2))
-# res[12].numpy()
-# self.get_last(inputs).numpy()
-# import tensorflow as tf
-# tf.random.uniform(shape=shape[0:1], maxval=255, dtype=tf.int32)
-# from ml_utils import print_error
-# ims(mask[0].numpy())
-# print_error(lambda :ims(mask[0]))
-# from models_utils import ops as K
-
-
-# self.model.layers[0](x)
-from raven_utils.models.loss import VTRavenLoss, SingleVTRavenLoss, create_uniform_mask
-
-
-def get_rav_trans(
-        data,
-        type_=9,
-        no=4,
-        extractor="ef",
-        loss_mode=create_uniform_mask,
-        output_size=10,
-        loss_weight=1.0,
-        out_layers=(1000, 1000, 1000),
-        pos_emd="cat",
-        base="seq",
-        inverse_image=True,
-        last="left",
-        epsilon="greedy",
-        epsilon_step=500,
-        mask_fn=None,
-        model=None,
-        loss="multi",
-        **kwargs):
-    if last == "left":
-        last = take_left
-    elif last == "mix":
-        last = mix
-    elif last == "empty":
-        last = empty_last
-    elif last == "start":
-        last = Sequential([Lambda(empty_last), BatchInitialWeight(initializer=init_weights)])
-
-    if epsilon == "greedy":
-        epsilon = EpsilonGreedy(step=epsilon_step)
-    elif epsilon == "soft":
-        epsilon = EpsilonSoft(step=epsilon_step)
-    elif epsilon is False:
-        epsilon = None
-
-    if epsilon:
-        trans_raven = TransRavenwithStep(
-            type_=type_,
-            no=no,
-            extractor=extractor,
-            output_size=output_size,
-            out_layer=out_layers,
-            pos_emd=pos_emd,
-            base=base,
-            last=last,
-            epsilon=epsilon,
-            **kwargs
-        )
-        return StepDict(bt(
-            DictModel(
-                Sequential([Lambda(lambda x: (255 - x[0], x[1])), trans_raven]) if inverse_image else trans_raven,
-                in_=[INPUTS, "step"],
-                name="Body"
-            ),
-            loss=VTRavenLoss(mode=loss_mode, classification=True, lw=(loss_weight, 1.0)),
-            loss_wrap=False),
-            add_step=epsilon_step,
-        )
-
-    trans_raven = img_sec_trans(
-        type_=type_,
-        no=no,
-        extractor=extractor,
-        model=model,
-        output_size=output_size,
-        out_layer=out_layers,
-        pos_emd=pos_emd,
-        base=base,
-        last=last,
-        epsilon=epsilon,
-        mask_fn=mask_fn,
-        **kwargs
-    )
-    if loss == "single":
-        loss = SingleVTRavenLoss
-    else:
-        loss = VTRavenLoss
-
-    # return bt(
-    #     DictModel(
-    #         Sequential([Lambda(lambda x: 255 - x), trans_raven]) if inverse_image else trans_raven,
-    #         inputs=INPUTS,
-    #         name="Body"
-    #     ),
-    #     loss=loss(mode=loss_mode, classification=True, lw=(loss_weight, 1.0)),
-    #     loss_wrap=False
-    # )
-
-    return bt([
-        DictModel(
-            Sequential([Lambda(lambda x: 255 - x), trans_raven]) if inverse_image else trans_raven,
-            in_=INPUTS,
-            name="Body"
-        ),
-
-    ],
-        loss=loss(mode=loss_mode, classification=True, lw=(loss_weight, 1.0)),
-        loss_wrap=False
-    )

raven_utils/models/transformer_2.py

@@ -1,146 +0,0 @@
-from functools import partial
-
-import tensorflow as tf
-from tensorflow.keras.layers import Lambda
-from tensorflow.python.keras import Sequential
-from models_utils import ops as K, SubClassing
-from models_utils.models.transformer import aug
-
-# from models_utils.models.loss import VTRavenLoss, create_uniform_mask, SingleVTRavenLoss
-from data_utils import DataGenerator, LOSS, TARGET, IMAGES
-from models_utils import DictModel, bt, INPUTS, BatchInitialWeight, build_functional_model, get_input_layer
-import models_utils.ops as K
-from models_utils.models.transformer.img_seq import init_weights, take_left, mix, empty_last
-from models_utils.models.transformer.img_seq2 import init_weights, take_left, mix, empty_last, img_sec_trans
-from models_utils.ops_core import IndexReshape
-from models_utils.random_ import EpsilonGreedy, EpsilonSoft
-from models_utils.step import StepDict
-
-from models_utils.models.transformer import aug
-
-# res = (1 - mask) * inputs[..., :self.last_index] + mask * tf.tile(self.get_last(inputs),
-#                                                                            (1, 1, 1, self.last_index))
-
-
-# from data_utils import ims
-# for i in range(50):
-#     ims(res[i].numpy().swapaxes(0, 2))
-# res[12].numpy()
-# self.get_last(inputs).numpy()
-# import tensorflow as tf
-# tf.random.uniform(shape=shape[0:1], maxval=255, dtype=tf.int32)
-# from ml_utils import print_error
-# ims(mask[0].numpy())
-# print_error(lambda :ims(mask[0]))
-# from models_utils import ops as K
-
-
-# self.model.layers[0](x)
-from raven_utils.constant import INDEX, LABELS
-from raven_utils.models.loss import VTRavenLoss, SingleVTRavenLoss, create_uniform_mask
-
-
-def get_matrix(inputs, index):
-    return tf.concat([inputs[:, :8], K.gather(inputs, index[:, 0])[:, None]], axis=1)
-
-
-def get_images(inputs):
-    return get_matrix(inputs[0], inputs[1])
-
-
-def random_last(inputs, max_=8):
-    index = K.init.label(max=max_, shape=[tf.shape(inputs[0])[0]])[..., None]
-    return get_matrix(inputs[0], index)
-
-
-def get_images_no_answer(inputs):
-    return inputs[0][:, :9]
-
-
-def repeat_last(inputs):
-    return inputs[0][:, list(range(8)) + [7]]
-
-
-def get_rav_trans(
-        data,
-        inverse_image=True,
-        loss_mode=create_uniform_mask,
-        loss_weight=1.0,
-        loss="multi",
-        number_loss=False,
-        plw=None,
-        pre="auto",
-        augmentation=None,
-        **kwargs):
-    if isinstance(data, DataGenerator):
-        data = data[0]['inputs'], data[0]['index']
-    # u = img_sec_trans(**kwargs)(get_images(data) if kwargs['mask'] == "random" else get_images_no_answer(data))
-    # u.shape
-    from keras import Model
-    if pre == "auto":
-        pre = get_images if kwargs['mask'] == "random" else get_images_no_answer
-    elif pre == "no_answer":
-        pre = get_images_no_answer
-    elif pre == "last":
-        pre = repeat_last
-    elif pre == "images":
-        pre = get_images
-    elif pre == "random_last":
-        pre = random_last
-    elif pre == "noise":
-        pre = SubClassing([get_matrix, partial(aug.noise, max_=8)])
-    elif pre == "batch_noise":
-        pre = SubClassing([get_matrix, partial(aug.batch_noise, max_=8)])
-
-    if augmentation == "transpose":
-        augmentation = aug.Transpose(axis=(0, 2, 1))
-        augmentation_label = aug.Transpose(axis=(0, 2, 1))
-    elif augmentation == "shuffle_col":
-        augmentation = aug.shuffle_col
-        augmentation_label = aug.shuffle_col
-    elif augmentation == "shuffle":
-        augmentation = aug.shuffle
-        augmentation_label = aug.shuffle
-    if augmentation:
-        augmentation = [
-            # DictModel(augmentation, IMAGES, IMAGES),
-            # DictModel(aug.reshape_static(pre(data),augmentation), IMAGES, IMAGES),
-            DictModel(aug.ReshapeStatic(augmentation), IMAGES, IMAGES),
-            DictModel(
-                aug.PartialModel(
-                    aug.ReshapeStatic(augmentation_label),
-                    last_axis=9)
-                , LABELS, LABELS)
-        ]
-    else:
-        augmentation = []
-
-    trans_raven = build_functional_model(
-        img_sec_trans(**kwargs),
-        # get_images(data) if kwargs['mask'] == "random" else get_images_no_answer(data)
-        pre(data)
-        # data[0]
-    )
-    if loss == "single":
-        loss = SingleVTRavenLoss
-    else:
-        loss = VTRavenLoss
-    if isinstance(loss_weight, float):
-        loss_weight = (loss_weight, 1.0)
-
-    return bt([
-        # DictModel(get_images if kwargs['mask'] == "random" else get_images_no_answer, [INPUTS, INDEX], IMAGES),
-        DictModel(pre, [INPUTS, INDEX], IMAGES),
-        *augmentation,
-        DictModel(
-            Sequential([Lambda(lambda x: 255 - x), trans_raven]) if inverse_image else trans_raven,
-            in_=IMAGES,
-            # inputs=INPUTS,
-            name="Body"
-        ),
-
-    ],
-        loss=loss(mode=loss_mode, classification=True, number_loss=number_loss, lw=loss_weight, plw=plw),
-        predict=LOSS,
-        loss_wrap=False
-    )

raven_utils/models/transformer_3.py

@@ -1,206 +0,0 @@
-import logging
-
-from loguru import logger
-from tensorflow.keras.layers import Lambda
-from tensorflow.keras.layers import Activation
-
-from grid_transformer import aug_trans
-from raven_utils.models.loss_3 import VTRavenLoss, SingleVTRavenLoss, create_uniform_mask
-from data_utils import get_shape, TakeDict
-
-from data_utils import DataGenerator, LOSS, TARGET, IMAGES
-from models_utils import DictModel, bt, INPUTS, BatchInitialWeight, build_functional, get_input_layer, Last, bm, \
-    add_end, AUGMENTATION
-# from report.select_ import SelectModel2, SelectModel, SelectModel9
-from experiment_utils.keras_model import load_weights as model_load_weights
-
-
-def get_rav_trans(
-        data,
-        loss_mode=create_uniform_mask,
-        loss_weight=2.0,
-        number_loss=False,
-        dry_run="auto",
-        plw=None,
-        **kwargs):
-    if isinstance(loss_weight, float):
-        loss_weight = (loss_weight, 1.0)
-
-    # seq_trans(**kwargs)(data[0])
-    # trans_raven = build_functional_model2(
-    #     seq_trans(**kwargs),
-    #     data[0],
-    #     batch=None
-    # )
-    trans_raven = build_functional(
-        model=aug_trans,
-        inputs_=data[0] if isinstance(data, DataGenerator) else data,
-        batch_=None,
-        dry_run=dry_run,
-        **kwargs
-    )
-
-    return bt(
-        model=trans_raven,
-        loss=VTRavenLoss(mode=loss_mode, classification=True, number_loss=number_loss, lw=loss_weight, plw=plw),
-        model_wrap=False,
-        predict=LOSS,
-        loss_wrap=False
-    )
-
-
-def rav_select_model(
-        data,
-        load_weights=None,
-        loss_weight=(0.01, 0.0),
-        plw=5.0,
-        result_metric="sparse_categorical_accuracy",
-        select_type=2,
-        select_out=0,
-        additional_out=0,
-        additional_copy=True,
-        tail_out=(1000, 1000),
-        **kwargs
-):
-    out_layers = Last()
-    if additional_out > 0:
-        model3 = get_rav_trans(
-            data,
-            plw=plw,
-            loss_weight=loss_weight,
-            **kwargs
-        )
-
-        model_load_weights(
-            model3,
-            load_weights,
-            # sample_data,
-            None,
-            template="weights_{epoch:02d}-{val_loss:.2f}",
-            key=result_metric,
-        )
-
-        if AUGMENTATION in kwargs and kwargs[AUGMENTATION] is not None:
-            index = -1
-        else:
-            index = -2
-
-        out = model3[0, index, :additional_out]
-        logger.info(f"Additional out from: {model3[0, index]}.")
-
-        if additional_out > 2:
-            out += [Activation("gelu")]
-        out_layers = bm([out_layers] + out, add_flatten=False)
-    model = get_rav_trans(
-        TakeDict(data[0])[:, 8:],
-        plw=plw,
-        loss_weight=loss_weight,
-        **{
-            **kwargs,
-            "out_layers": out_layers,
-        }
-        # **{**as_dict(p.mp), "show_shape": True, "save_shape": f"output/shapes/type_{p.mp.type_}.json"},
-    )
-    # from data_utils.ops import Equal
-    # o = []
-    # for i in range(1, 3):
-    #     for j in range(2):
-    #         o.append(
-    #             Equal(
-    #                 # model[0,:,-2, i].variables[j],
-    #                 model2[0, :, -2, i].variables[j],
-    #                 # out_layers[i].variables[j]
-    #                 second_pooling[i].variables[j]
-    #             ).equal
-    #         )
-    # assert all(o)
-    # model = get_rav_trans(
-    #     # TakeDict(val_generator[0])[:, 8:],
-    #     # TakeDict(val_generator[0])[:, 8:],
-    #     val_generator[0],
-    #     plw=p.plw,
-    #     loss_weight=p.loss_weight,
-    #     **{**as_dict(p.mp),
-    #        # "out_layers": out_layers,
-    #        }
-    #     # **{**as_dict(p.mp), "show_shape": True, "save_shape": f"output/shapes/type_{p.mp.type_}.json"},
-    # )
-    model_load_weights(model,
-                       load_weights,
-                       # sample_data,
-                       None,
-                       template="weights_{epoch:02d}-{val_loss:.2f}",
-                       key=result_metric,
-                       )
-    # model.compile()
-    # model.evaluate(val_generator.data[:1000])
-    # model(TakeDict(val_generator[0])[:, 8:])
-    trans_raven = model[0]
-    # s = trans_raven(TakeDict(val_generator[0])[:, 8:])
-    if select_type == 2:
-        second_pooling = Lambda(lambda x: x[:, :-1])
-    else:
-        second_pooling = Last()
-    if additional_out > 0:
-        if additional_copy:
-            model4 = get_rav_trans(
-                data,
-                plw=plw,
-                loss_weight=loss_weight,
-                **kwargs
-            )
-            model_load_weights(model4,
-                               load_weights,
-                               # sample_data,
-                               None,
-                               template="weights_{epoch:02d}-{val_loss:.2f}",
-                               key=result_metric,
-                               )
-
-            if AUGMENTATION in kwargs and kwargs[AUGMENTATION] is not None:
-                index = -1
-            else:
-                index = -2
-            out2 = model4[0, index, :additional_out]
-            logger.info(f"Additional out from: {model4[0, index]}.")
-
-            if additional_out > 2:
-                out2 += [Activation("gelu")]
-        else:
-            out2 = out
-
-        second_pooling = bm([second_pooling] + out2, add_flatten=False)
-
-    model2 = get_rav_trans(
-        TakeDict(data[0])[:, 8:],
-        plw=plw,
-        loss_weight=loss_weight,
-        **{
-            **kwargs,
-            "out_layers": second_pooling,
-        }
-        # **{**as_dict(p.mp), "show_shape": True, "save_shape": f"output/shapes/type_{p.mp.type_}.json"},
-    )
-    model_load_weights(
-        model2,
-        load_weights,
-        # sample_data,
-        None,
-        template="weights_{epoch:02d}-{val_loss:.2f}",
-        key=result_metric,
-    )
-    if select_type == 0:
-        # not working
-        trans_raven2 = model2[0]
-    else:
-        trans_raven2 = model2[0]
-    tail = add_end(out_layers=tail_out, output_size=8 if select_out else 1)
-    # trans_raven2.mask_fn = ImageMask(last=take_by_index)
-    if select_type == 2:
-        select_model_class = SelectModel2
-    elif select_type == 1:
-        select_model_class = SelectModel
-    else:
-        select_model_class = SelectModel9
-    select_model = select_model_class(trans_raven, model2=trans_raven2, tail=tail, select_out=select_out)
-    return select_model

raven_utils/models/uitls_.py

@@ -1,16 +0,0 @@
-import tensorflow as tf
-import tensorflow.experimental.numpy as tnp
-from tensorflow.keras import Model
-import raven_utils as rv
-
-
-class RangeMask(Model):
-    def __init__(self):
-        super().__init__()
-        ranges = tf.tile(tf.range(rv.entity.INDEX[-1])[None], [rv.group.NO, 1])
-        start_index = rv.entity.INDEX[:-1][:, None]
-        end_index = rv.entity.INDEX[1:][:, None]
-        self.mask = tnp.array((start_index <= ranges) & (ranges < end_index))
-
-    def call(self, inputs):
-        return self.mask[inputs]