回答編集履歴
2
修正
test
CHANGED
@@ -16,21 +16,19 @@
|
|
16
16
|
|
17
17
|
self.backbone_type = backbone_type
|
18
18
|
|
19
|
-
self.
|
19
|
+
if self.backbone_type == "ResNet50":
|
20
20
|
|
21
|
-
models.
|
21
|
+
self.model = models._utils.IntermediateLayerGetter(
|
22
22
|
|
23
|
+
models.resnet50(), {"layer1": 2, "layer2": 3, "layer3": 4}
|
24
|
+
|
23
|
-
)
|
25
|
+
)
|
24
26
|
|
25
27
|
|
26
28
|
|
27
29
|
def forward(self, x):
|
28
30
|
|
29
|
-
if self.backbone_type == "ResNet50":
|
30
|
-
|
31
|
-
|
32
|
-
|
33
|
-
|
31
|
+
return self.model(x)
|
34
32
|
|
35
33
|
```
|
36
34
|
|
1
修正
test
CHANGED
@@ -143,3 +143,11 @@
|
|
143
143
|
to("cuda") や cuda() は nn.Module 内の子の nn.Module を GPU に転送するというものですが、そもそも子として認識されていない nn.Module は転送されないままということになります。
|
144
144
|
|
145
145
|
質問のコードだと models._utils.IntermediateLayerGetter() の返り値を `__init__()` 内で属性として定義していなかったので、子として認識されておらず、`cuda()` を実行しても GPU に転送されないままとなっています。
|
146
|
+
|
147
|
+
|
148
|
+
|
149
|
+
コードを読むと、nn.module 周りがどのように実現されているかがわかります。
|
150
|
+
|
151
|
+
|
152
|
+
|
153
|
+
[torch.nn.modules.module — PyTorch 1.6.0 documentation](https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module.children)
|