Bug
Traceback (most recent call last):
  File "/Users/simon/.pyenv/versions/3.7.4/lib/python3.7/site-packages/torch/jit/__init__.py", line 1034, in trace_module
    module._c._create_method_from_trace(method_name, func, example_inputs, var_lookup_fn, _force_outplace)
RuntimeError: type INTERNAL ASSERT FAILED at ../torch/csrc/jit/ir.h:1266, please report a bug to PyTorch.  (setType at ../torch/csrc/jit/ir.h:1266)
frame #0: c10::Error::Error(c10::SourceLocation, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&) + 135 (0x11ac06787 in libc10.dylib)
frame #1: torch::jit::Value::setType(std::__1::shared_ptr<c10::Type>) + 459 (0x11eee6eab in libtorch.dylib)
frame #2: torch::jit::Graph::createGetAttr(torch::jit::Value*, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&) + 259 (0x11ef6ee03 in libtorch.dylib)
frame #3: torch::jit::(anonymous namespace)::ConvertTracedAttrReferences::convertAttrReferencesToLocalGetAttrs(torch::jit::Block*, c10::QualifiedName const&, torch::jit::Value*) + 714 (0x11efdc3da in libtorch.dylib)
frame #4: torch::jit::FixupTraceScopeBlocks(std::__1::shared_ptr<torch::jit::Graph>&, torch::jit::script::Module*) + 316 (0x11efd935c in libtorch.dylib)
frame #5: torch::jit::tracer::trace(std::__1::vector<c10::IValue, std::__1::allocator<c10::IValue> >, std::__1::function<std::__1::vector<c10::IValue, std::__1::allocator<c10::IValue> > (std::__1::vector<c10::IValue, std::__1::allocator<c10::IValue> >)> const&, std::__1::function<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > (at::Tensor const&)>, bool, torch::jit::script::Module*) + 2127 (0x11f1c8fef in libtorch.dylib)
frame #6: torch::jit::tracer::createGraphByTracing(pybind11::function const&, std::__1::vector<c10::IValue, std::__1::allocator<c10::IValue> >, pybind11::function const&, bool, torch::jit::script::Module*) + 361 (0x11b7d3a49 in libtorch_python.dylib)
frame #7: void pybind11::cpp_function::initialize<torch::jit::script::initJitScriptBindings(_object*)::$_13, void, torch::jit::script::Module&, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&, pybind11::function, pybind11::tuple, pybind11::function, bool, pybind11::name, pybind11::is_method, pybind11::sibling>(torch::jit::script::initJitScriptBindings(_object*)::$_13&&, void (*)(torch::jit::script::Module&, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&, pybind11::function, pybind11::tuple, pybind11::function, bool), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&)::'lambda'(pybind11::detail::function_call&)::__invoke(pybind11::detail::function_call&) + 319 (0x11b814d9f in libtorch_python.dylib)
frame #8: pybind11::cpp_function::dispatcher(_object*, _object*, _object*) + 3372 (0x11b1479fc in libtorch_python.dylib)
<omitting python frames>
frame #13: __pyx_f_18_pydevd_frame_eval_35pydevd_frame_evaluator_darwin_37_64_get_bytecode_while_frame_eval + 6471 (0x11aa0cef7 in pydevd_frame_evaluator_darwin_37_64.cpython-37m-darwin.so)
frame #18: __pyx_f_18_pydevd_frame_eval_35pydevd_frame_evaluator_darwin_37_64_get_bytecode_while_frame_eval + 6471 (0x11aa0cef7 in pydevd_frame_evaluator_darwin_37_64.cpython-37m-darwin.so)
frame #23: __pyx_f_18_pydevd_frame_eval_35pydevd_frame_evaluator_darwin_37_64_get_bytecode_while_frame_eval + 6471 (0x11aa0cef7 in pydevd_frame_evaluator_darwin_37_64.cpython-37m-darwin.so)
frame #27: __pyx_f_18_pydevd_frame_eval_35pydevd_frame_evaluator_darwin_37_64_get_bytecode_while_frame_eval + 6471 (0x11aa0cef7 in pydevd_frame_evaluator_darwin_37_64.cpython-37m-darwin.so)
frame #35: __pyx_f_18_pydevd_frame_eval_35pydevd_frame_evaluator_darwin_37_64_get_bytecode_while_frame_eval + 6471 (0x11aa0cef7 in pydevd_frame_evaluator_darwin_37_64.cpython-37m-darwin.so)
frame #40: __pyx_f_18_pydevd_frame_eval_35pydevd_frame_evaluator_darwin_37_64_get_bytecode_while_frame_eval + 6471 (0x11aa0cef7 in pydevd_frame_evaluator_darwin_37_64.cpython-37m-darwin.so)
frame #57: start + 1 (0x7fff6bfa93d5 in libdyld.dylib)
frame #58: 0x0 + 10 (0xa in ???)

Environment
PyTorch version: 1.4.0
Is debug build: No
CUDA used to build PyTorch: None

OS: Mac OSX 10.14.6
GCC version: Could not collect
CMake version: Could not collect

Python version: 3.7
Is CUDA available: No
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA

Versions of relevant libraries:
[pip3] numpy==1.18.2
[pip3] torch==1.4.0
[pip3] torchvision==0.5.0
[conda] Could not collect

Repeat
from collections import defaultdict

import torch
import torch.jit
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models

from utils import *

class MyConv2D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
        super(MyConv2D, self).__init__()
        self.weight = torch.zeros((out_channels, in_channels, kernel_size, kernel_size))
        self.bias = torch.zeros(out_channels)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = (kernel_size, kernel_size)
        self.stride = (stride, stride)

    def forward(self, x):
        return F.conv2d(x, self.weight, self.bias, self.stride)

    def extra_repr(self):
        s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}'
             ', stride={stride}')
        return s.format(**self.__dict__)


class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv = nn.Sequential(
            *ConvLayer(channels, channels, kernel_size=3, stride=1),
            *ConvLayer(channels, channels, kernel_size=3, stride=1, relu=False)
        )

    def forward(self, x):
        return self.conv(x) + x


def ConvLayer(in_channels, out_channels, kernel_size=3, stride=1,
              upsample=None, instance_norm=True, relu=True, trainable=False):
    layers = []
    if upsample:
        layers.append(nn.Upsample(mode='nearest', scale_factor=upsample))
    layers.append(nn.ReflectionPad2d(kernel_size // 2))
    if trainable:
        layers.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride))
    else:
        layers.append(MyConv2D(in_channels, out_channels, kernel_size, stride))
    if instance_norm:
        layers.append(nn.InstanceNorm2d(out_channels))
    if relu:
        layers.append(nn.ReLU())
    return layers


class TransformNet(nn.Module):
    def __init__(self, base=8):
        super(TransformNet, self).__init__()
        self.base = base
        self.weights = []
        self.downsampling = nn.Sequential(
            *ConvLayer(3, base, kernel_size=9, trainable=True),
            *ConvLayer(base, base * 2, kernel_size=3, stride=2),
            *ConvLayer(base * 2, base * 4, kernel_size=3, stride=2),
        )
        self.residuals = nn.Sequential(*[ResidualBlock(base * 4) for _ in range(5)])
        self.upsampling = nn.Sequential(
            *ConvLayer(base * 4, base * 2, kernel_size=3, upsample=2),
            *ConvLayer(base * 2, base, kernel_size=3, upsample=2),
            *ConvLayer(base, 3, kernel_size=9, instance_norm=False, relu=False, trainable=True),
        )
        self.get_param_dict()

    def forward(self, X):
        y = self.downsampling(X)
        y = self.residuals(y)
        y = self.upsampling(y)
        return y

    def get_param_dict(self):
        """找出该网络所有 MyConv2D 层，计算它们需要的权值数量"""
        param_dict = defaultdict(int)
        for value in self.named_modules():
            if isinstance(value[1], MyConv2D):
                param_dict[value[0]] += int(np.prod(value[1].weight.shape))
                param_dict[value[0]] += int(np.prod(value[1].bias.shape))
        return param_dict

    def set_my_attr(self, name, value):
        # 下面这个循环是一步步遍历类似 residuals.0.conv.1 的字符串，找到相应的权值
        target = self
        for x in name.split('.'):
            if x.isnumeric():
                target = target.__getitem__(int(x))
            else:
                target = getattr(target, x)

        # 设置对应的权值
        n_weight = np.prod(target.weight.shape)
        target.weight = value[:n_weight].view(target.weight.shape)
        target.bias = value[n_weight:].view(target.bias.shape)

    def set_weights(self, weights, i=0):
        """输入权值字典，对该网络所有的 MyConv2D 层设置权值"""
        for name, param in weights.items():
            self.set_my_attr(name, weights[name][i])


class MetaNet(nn.Module):
    def __init__(self, param_dict):
        super(MetaNet, self).__init__()
        self.hidden = nn.Linear(1920, 128 * len(param_dict))
        self.fc_list = []
        for index, (name, params) in enumerate(param_dict.items()):
            setattr(self, name, nn.Linear(128, params))
            self.fc_list.append(name)

    def forward(self, mean_std_features):
        hidden = F.relu(self.hidden(mean_std_features))
        filters = []
        labels = []
        for index, name in enumerate(self.fc_list):
            labels.append(torch.ones(1, dtype=torch.int8) * index)
            filters.append(getattr(self, name)(hidden[:, index * 128:(index + 1) * 128]))
        return tuple(filters)

def test_mobile():
    transform_net = TransformNet(base=32).cpu().eval()
    meta_net = MetaNet(transform_net.get_param_dict()).cpu().eval()
    input_1 = torch.ones(1, 1920)
    meta_net_mobile = torch.jit.trace(meta_net, input_1)
    meta_net_mobile.save("models/meta_net_mobile.pt")
    print(meta_net_mobile.code)

if __name__ == "__main__":
    test_mobile()