bes-dev
commented
1 year ago If I try to use the same script but with FC model from AITemplate docs, it works well:
import logging
from aitemplate.compiler import ops
from aitemplate.frontend import nn
# AIT utils
from aitemplate.compiler import compile_model
from aitemplate.frontend import Tensor
from aitemplate.testing import detect_target
class AITSimpleModel(nn.Module):
def __init__(self, hidden, eps: float = 1e-5):
super().__init__()
self.dense1 = nn.Linear(hidden, 4 * hidden, specialization="fast_gelu")
self.dense2 = nn.Linear(4 * hidden, hidden)
self.layernorm = nn.LayerNorm(hidden, eps=eps)
def forward(self, input):
hidden_states = self.dense1(input)
hidden_states = self.dense2(hidden_states)
hidden_states = hidden_states + input
hidden_states = self.layernorm(hidden_states)
return hidden_states
def mark_output(y):
if type(y) is not tuple:
y = (y,)
for i in range(len(y)):
y[i]._attrs["is_output"] = True
y[i]._attrs["name"] = "output_%d" % (i)
y_shape = [d._attrs["values"][0] for d in y[i]._attrs["shape"]]
print("AIT output_{} shape: {}".format(i, y_shape))
def compile_moc(
batch_size,
input_size,
c_in = 3,
c_out = 8,
use_fp16_acc = False,
convert_conv_to_gemm = False,
output_dir="./tmp/"
):
ait_model = AITSimpleModel(32)
ait_input = Tensor(
shape=[1, 32],
name="input0",
is_input=True,
)
ait_model.name_parameter_tensor()
ait_out = ait_model(ait_input)
mark_output(ait_out)
target = detect_target(
use_fp16_acc=use_fp16_acc,
convert_conv_to_gemm=convert_conv_to_gemm
)
compile_model(
ait_out,
target,
output_dir,
"moc"
)
def main():
logging.getLogger().setLevel(logging.INFO)
logger = logging.getLogger()
logger.info("Compile model...")
compile_moc(
batch_size=1,
input_size=256,
use_fp16_acc=True,
convert_conv_to_gemm=True,
)
if __name__ == "__main__":
main()
hl475
aakhundov
fatihcelikbas
j3698w
delldu
When I try to compile simple convolution network, compilation process crash because conv2d.attrs_['op_instance'] is empty for the convolution layer. How can I fix it?
Behaviour can be reproduced by this script:
Output:
conv2d.attrs_ state: