fffiloni commited on
Commit
c381dd3
1 Parent(s): 5a972ab

Create model.py

Browse files
Files changed (1) hide show
  1. utils/model.py +32 -0
utils/model.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ import os
3
+ import time
4
+ import pickle
5
+
6
+ import torch
7
+ import torch.distributed as dist
8
+
9
+ from fvcore.nn import FlopCountAnalysis
10
+ from fvcore.nn import flop_count_table
11
+ from fvcore.nn import flop_count_str
12
+
13
+ logger = logging.getLogger(__name__)
14
+
15
+
16
+ NORM_MODULES = [
17
+ torch.nn.BatchNorm1d,
18
+ torch.nn.BatchNorm2d,
19
+ torch.nn.BatchNorm3d,
20
+ torch.nn.SyncBatchNorm,
21
+ # NaiveSyncBatchNorm inherits from BatchNorm2d
22
+ torch.nn.GroupNorm,
23
+ torch.nn.InstanceNorm1d,
24
+ torch.nn.InstanceNorm2d,
25
+ torch.nn.InstanceNorm3d,
26
+ torch.nn.LayerNorm,
27
+ torch.nn.LocalResponseNorm,
28
+ ]
29
+
30
+ def register_norm_module(cls):
31
+ NORM_MODULES.append(cls)
32
+ return cls