FilesExpand file tree

master

/

modelutils.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

16 lines (12 loc) · 397 Bytes

master

/

modelutils.py

Copy path

Top

File metadata and controls

• Code
• Blame

16 lines (12 loc) · 397 Bytes

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

import torch

import torch.nn as nn

DEV = torch.device('cuda:0')

def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''):

if type(module) in layers:

return {name: module}

res = {}

for name1, child in module.named_children():

res.update(find_layers(

child, layers=layers, name=name + '.' + name1 if name != '' else name1

))

return res