def convert_graph_layout(mod, desired_layout): """Alter the layout of the input graph. Parameters ---------- mod : tvm.relay.Module The relay module to convert. desired_layout : str The layout to convert to. Returns ------- mod : tvm.relay.Module The converted module. """ # Assume for the time being that graphs only have # conv2d as heavily-sensitive operators. desired_layouts = { "nn.conv2d": [desired_layout, "default"], "qnn.conv2d": [desired_layout, "default"], } # Convert the layout of the graph where possible. seq = transform.Sequential([ relay.transform.RemoveUnusedFunctions(), relay.transform.ConvertLayout(desired_layouts), ]) with transform.PassContext(opt_level=3): try: return seq(mod) except Exception as err: raise TVMCException("Error converting layout to {0}: {1}".format( desired_layout, str(err)))
return op def visit_var(self, var: relay.Var): new_ty = self.ty_mut.visit(var.checked_type) return relay.Var(name_hint=var.name_hint, type_annotation=new_ty) class TensorDTypeMutator(relay.TypeMutator): def __init__(self, tgt_ty: str): super().__init__() self.tgt_ty = tgt_ty def visit_tensor_type(self, tt: relay.TensorType): return relay.TensorType(tt.concrete_shape, dtype=self.tgt_ty) inputs = keras.Input(shape=(224, 224, 3), batch_size=4) x = layers.Conv2D(64, 3, use_bias=False, padding='same')(inputs) x = layers.BatchNormalization(epsilon=1e-5)(x) x = layers.GlobalAvgPool2D()(x) keras_model = keras.Model(inputs=inputs, outputs=x) keras_model.summary() ir_mod, params = relay.frontend.from_keras(keras_model, shape={'input_1': (4, 3, 224, 224)}) print(ir_mod) ir_mod = transform.Sequential(passes=[ AlterDType('float16'), relay.transform.InferType(), ])(ir_mod) print(ir_mod)