提交 132ec21f 编写于 作者: W weishengyu

reverse res_dict

上级 56c513fd
......@@ -16,10 +16,6 @@ class TheseusLayer(nn.Layer):
super(TheseusLayer, self).__init__()
self.res_dict = None
def forward(self, *input, res_dict=None, **kwargs):
if res_dict is not None:
self.res_dict = res_dict
# stop doesn't work when stop layer has a parallel branch.
def stop_after(self, stop_layer_name: str):
after_stop = False
......@@ -44,6 +40,7 @@ class TheseusLayer(nn.Layer):
if isinstance(self._sub_layers[layer_i], TheseusLayer):
self._sub_layers[layer_i].register_forward_post_hook(
self._sub_layers[layer_i]._save_sub_res_hook)
self._sub_layers[layer_i].res_dict = self.res_dict
if isinstance(self._sub_layers[layer_i], TheseusLayer):
self._sub_layers[layer_i]._update_res(return_patterns)
......
......@@ -84,8 +84,7 @@ class ConvBlock(TheseusLayer):
self.max_pool = MaxPool2D(kernel_size=2, stride=2, padding=0)
self.relu = nn.ReLU()
def forward(self, inputs, res_dict=None):
super(ConvBlock, self).forward(inputs, res_dict=res_dict)
def forward(self, inputs):
x = self.conv1(inputs)
x = self.relu(x)
if self.groups == 2 or self.groups == 3 or self.groups == 4:
......@@ -141,7 +140,6 @@ class VGGNet(TheseusLayer):
self._update_res(return_patterns)
def forward(self, inputs, res_dict=None):
super(VGGNet, self).forward(inputs, res_dict=res_dict)
x = self.conv_block_1(inputs, res_dict)
x = self.conv_block_2(x, res_dict)
x = self.conv_block_3(x, res_dict)
......@@ -155,6 +153,9 @@ class VGGNet(TheseusLayer):
x = self.relu(x)
x = self.drop(x)
x = self.fc3(x)
if self.res_dict:
for res_key in self.res_dict:
res_dict[res_key] = self.res_dict.pop(res_key)
return x
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册