model_interface.cu:231: Error: Constant pretrained_model_patch_embed_proj_weight was not set! Set the value with set_constant.

[ADongGu]

Hi.
I have some problems, i'm doing a simple remodel, and I get this kind of error. It's kind of hard to understand

import torch

from aitemplate.compiler import compile_model
from aitemplate.frontend import nn, Tensor
from aitemplate.testing import detect_target
from aitemplate.compiler.ops import flatten
from aitemplate.frontend import IntVar, Tensor

class PTSimpleModel(torch.nn.Module):
    def __init__(self, img_size=512, patch_size=16, in_chans=3, embed_dim=1024, norm_layer=None, flatten=True):
        super().__init__()
        self.proj = torch.nn.Conv2d(in_chans, embed_dim, kernel_size=(16,16), stride=(16,16))

    def forward(self, x):
        x = self.proj(x)
        x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        return x

class AITSimpleModel(nn.Module):
    def __init__(self, img_size=512, patch_size=16, in_chans=3, embed_dim=1024, norm_layer=None, flatten=True):
        super().__init__()
        self.proj = nn.Conv2dBiasFewChannels(in_channels=in_chans,out_channels=embed_dim, kernel_size=16, stride=16,dtype="float32")

    def forward(self, x):
        x = self.proj(x)
        x = flatten(1,2)(x)
        return x

def map_pt_params(ait_model, pt_model):
    ait_model.name_parameter_tensor()
    pt_params = dict(pt_model.named_parameters())
    mapped_pt_params = {}
    for name, _ in ait_model.named_parameters():
        ait_name = name.replace(".", "_")
        assert name in pt_params
        mapped_pt_params[ait_name] = pt_params[name]
    return mapped_pt_params

def verify_simple_model(batch_size=1024, hidden=512):
    pt_model = PTSimpleModel().cuda()
    pt_model.eval()
    # x = torch.randn([1, 3, 672,512]).cuda()
    # y = pt_model(x)
    # print(y.shape)
    ait_model = AITSimpleModel()
    X = Tensor(
        shape=[1,672,512,3],
        name="X",
        dtype="float32",
        is_input=True,
    )
    Y = ait_model(X)
    Y._attrs["is_output"] = True
    Y._attrs["name"] = "Y"

    weights = map_pt_params(ait_model, pt_model)

    target = detect_target()
    with compile_model(
        Y, target, "./tmp", "simple_model_demo", constants=weights
    ) as module:
        ait_input = [torch.randn(1,672,512,3).cuda()]
        ys = []
        num_outputs = len(module.get_output_name_to_index_map())
        for i in range (num_outputs):
            shape = module.get_output_maximum_shape(i)
            ys.append(torch.empty(shape).cuda())
        module.run_with_tensors(ait_input,ys,graph_mode=False)

verify_simple_model()

[ADongGu]

re compile AITemplate