Print the str and parameter number of a network. Args: net (nn.Module)
(self, net)
| 146 | |
| 147 | @master_only |
| 148 | def print_network(self, net): |
| 149 | """Print the str and parameter number of a network. |
| 150 | |
| 151 | Args: |
| 152 | net (nn.Module) |
| 153 | """ |
| 154 | if isinstance(net, (DataParallel, DistributedDataParallel)): |
| 155 | net_cls_str = f'{net.__class__.__name__} - {net.module.__class__.__name__}' |
| 156 | else: |
| 157 | net_cls_str = f'{net.__class__.__name__}' |
| 158 | |
| 159 | net = self.get_bare_model(net) |
| 160 | # net_str = str(net) |
| 161 | net_params = sum(map(lambda x: x.numel(), net.parameters())) |
| 162 | |
| 163 | logger = get_root_logger() |
| 164 | logger.info(f'Network: {net_cls_str}, with parameters: {net_params:,d}') |
| 165 | # logger.info(net_str) |
| 166 | |
| 167 | def _set_lr(self, lr_groups_l): |
| 168 | """Set learning rate for warm-up. |
no test coverage detected