github/SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/SHARK-Studio.git synced 2026-01-10 06:17:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

Install torch version required by torch-mlir when setting up importer venv. (#486)

Browse Source
This commit is contained in:
Ean Garvey
2022-11-14 02:31:01 -06:00
committed by GitHub
parent d9f4a9954a
commit 45677c1e23
10 changed files with 15 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
setup_venv.sh
View File

@@ -105,8 +105,6 @@ else
echo "Not installing a backend, please make sure to add your backend to PYTHONPATH"
fi
$PYTHON -m pip install -e . -f https://llvm.github.io/torch-mlir/package-index/ -f ${RUNTIME} -f https://download.pytorch.org/whl/nightly/torch/
if [[ ! -z "${IMPORTER}" ]]; then
echo "${Yellow}Installing importer tools.."
if [[ $(uname -s) = 'Linux' ]]; then
@@ -119,6 +117,9 @@ if [[ ! -z "${IMPORTER}" ]]; then
$PYTHON -m pip install --no-warn-conflicts --upgrade -r "$TD/requirements-importer-macos.txt" -f ${RUNTIME} --extra-index-url https://download.pytorch.org/whl/nightly/cpu
fi
fi
$PYTHON -m pip install --no-warn-conflicts -e . -f https://llvm.github.io/torch-mlir/package-index/ -f ${RUNTIME} -f https://download.pytorch.org/whl/nightly/torch/
if [[ $(uname -s) = 'Linux' && ! -z "${BENCHMARK}" ]]; then
$PYTHON -m pip uninstall -y torch torchvision
$PYTHON -m pip install --pre torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cu116

2
shark/examples/shark_training/bert_training_tf.py
View File

@@ -53,7 +53,7 @@ class BertModule(tf.Module):
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
],
jit_compile=True
jit_compile=True,
)
def forward(self, inputs, labels):
with tf.GradientTape() as tape:

2
tank/examples/bert_fine_tuning/bert_fine_tune_tf.py
View File

@@ -73,7 +73,7 @@ class BertModule(tf.Module):
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
],
jit_compile=True
jit_compile=True,
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:

2
tank/examples/bert_tf/bert_large_gen.py
View File

@@ -61,7 +61,7 @@ class BertModule(tf.Module):
), # input2: segment_ids
tf.TensorSpec([BATCH_SIZE], tf.int32), # input3: labels
],
jit_compile=True
jit_compile=True,
)
def learn(self, input_word_ids, input_mask, segment_ids, labels):
with tf.GradientTape() as tape:

2
tank/examples/bert_tf/bert_large_run.py
View File

@@ -58,7 +58,7 @@ class BertModule(tf.Module):
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
],
jit_compile=True
jit_compile=True,
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:

2
tank/examples/bert_tf/bert_large_tf.py
View File

@@ -51,7 +51,7 @@ class BertModule(tf.Module):
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
],
jit_compile=True
jit_compile=True,
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:

2
tank/examples/bert_tf/bert_small_gen.py
View File

@@ -58,7 +58,7 @@ class BertModule(tf.Module):
), # input2: segment_ids
tf.TensorSpec([BATCH_SIZE], tf.int32), # input3: labels
],
jit_compile=True
jit_compile=True,
)
def learn(self, input_word_ids, input_mask, segment_ids, labels):
with tf.GradientTape() as tape:

2
tank/examples/bert_tf/bert_small_run.py
View File

@@ -54,7 +54,7 @@ class BertModule(tf.Module):
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
],
jit_compile=True
jit_compile=True,
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:

2
tank/examples/bert_tf/bert_small_tf_run.py
View File

@@ -47,7 +47,7 @@ class BertModule(tf.Module):
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
],
jit_compile=True
jit_compile=True,
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:

5
tank/model_utils_tf.py
View File

@@ -191,7 +191,10 @@ class ResNetModule(tf.Module):
self.m = tf_model
self.m.predict = lambda x: self.m.call(x, training=False)
@tf.function(input_signature=[tf.TensorSpec(INPUT_SHAPE, tf.float32)], jit_compile=True)
@tf.function(
input_signature=[tf.TensorSpec(INPUT_SHAPE, tf.float32)],
jit_compile=True,
)
def forward(self, inputs):
return self.m.predict(inputs)