|
- import os
- import logging
- import mindspore as ms
- import numpy as np
- from mindspore import nn, Tensor
- from mindspore import FixedLossScaleManager, Model
- from mindspore.communication import init, get_rank, get_group_size
-
- from mindcv.models import create_model
- from mindcv.data import create_dataset, create_transforms, create_loader
- from mindcv.loss import create_loss
- from mindcv.optim import create_optimizer
- from mindcv.scheduler import create_scheduler
- from mindcv.utils import StateMonitor, Allreduce
- from config import parse_args
-
-
- # Define Your Network
- logger = logging.getLogger('debug')
-
- network = create_model("googlenet")
- logger.info(network)
- id2name = {}
- def forward_hook_fn(cell_id, inputs, output):
- logger.info(f"[cell] id: {cell_id}, name: {id2name[cell_id]}")
- logger.info(f"[forward in] min: {inputs[0].min()}, max: {inputs[0].max()}, has_nan: {ms.ops.isnan(inputs[0]).any()}")
- logger.info(f"[forward out] min: {output.min()}, max: {output.max()}, has_nan: {ms.ops.isnan(output).any()}")
-
- def backward_hook_fn(cell_id, grad_inputs, grad_outputs):
- logger.info(f"[cell] id: {cell_id}, name: {id2name[cell_id]}")
- logger.info(f"[backward in] min: {grad_inputs[0].min()}, max: {grad_inputs[0].max()}, has_nan: {ms.ops.isnan(grad_inputs[0]).any()}")
- logger.info(f"[backward out] min: {grad_outputs[0].min()}, max: {grad_outputs[0].max()}, has_nan: {ms.ops.isnan(grad_outputs[0]).any()}")
-
- for name, cell in network.cells_and_names():
- id2name[cell.cls_name + "(" + str(id(cell)) + ")"] = name
- # Filter out layers based on cell type
- if isinstance(cell, (ms.nn.SequentialCell, ShuffleV1Block)):
- continue
- # Filter out layers based on cell name
- if "convs" in name:
- continue
- logger.info(f"register hook to {name}")
- cell.register_forward_hook(forward_hook_fn)
- cell.register_backward_hook(backward_hook_fn)
-
|