34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
from pathlib import Path
|
|
|
|
import tensorflow as tf
|
|
|
|
from src.common import DataType, Op
|
|
from src.tf_2.base import TFBase
|
|
|
|
|
|
class TFMatmulBench(TFBase):
|
|
def __init__(self, output_path: Path):
|
|
super().__init__(output_path, Op.MATMUL)
|
|
|
|
def experiment(self, experiment_args: tuple[int, int], length: int, dtype: tf.DType, device: tf.device):
|
|
shape_1, shape_2 = experiment_args
|
|
with device:
|
|
tensor_1 = tf.ones(shape_1, dtype=dtype)
|
|
tensor_2 = tf.ones(shape_2, dtype=dtype)
|
|
|
|
for _ in range(length):
|
|
_ = tensor_1 @ tensor_2
|
|
|
|
def name(self, experiment_args: tuple[int, int]) -> str:
|
|
shape_1, shape_2 = experiment_args
|
|
return f'{shape_1[0]}x{shape_1[1]} @ {shape_2[0]}x{shape_2[1]}'
|
|
|
|
def mop(self, experiment_args: tuple[int, int]) -> float:
|
|
shape_1, shape_2 = experiment_args
|
|
return (shape_1[0] * shape_2[1] / 1000_000) * 2 * (shape_1[1] - 1)
|
|
|
|
def run(self,
|
|
experiment_args: list[tuple[tuple[int, int], tuple[int, int]]],
|
|
experiment_count: int,
|
|
data_type: DataType):
|
|
super().run(experiment_args, experiment_count, data_type)
|