def log_layer_weights(self, tag, layer_names, step=None):
    step = (step or self.global_step)
    from arnet.modeling.models import MODEL_REGISTRY
    if (isinstance(self.model, MODEL_REGISTRY.get('CNN_Li2020')) or isinstance(self.model, MODEL_REGISTRY.get('SimpleC3D'))):
        for layer_name in layer_names:
            layer = utils.get_layer(self.model, layer_name)
            if isinstance(layer, torch.nn.Conv3d):
                fig = utils.draw_conv2d_weight(layer.weight)
                image_tensor = utils.fig2rgb(fig)
                save_name = (tag + f'/unscaled/{layer_name}')
                self.logger.experiment.add_image(save_name, image_tensor, step)
                save_name += f'/{step}.png'
                mlflow.log_figure(fig, save_name)
                fig = utils.draw_conv2d_weight(layer.weight, vmin=(- 0.3), vmax=0.3)
                image_tensor = utils.fig2rgb(fig)
                save_name = (tag + f'/uniform_scaled/{layer_name}')
                self.logger.experiment.add_image(save_name, image_tensor, step)
                save_name += f'/{step}.png'
                mlflow.log_figure(fig, save_name)
def log_layer_weights(self, tag, layer_names, step=None):
    step = (step or self.global_step)
    from arnet.modeling.models import MODEL_REGISTRY
    if (isinstance(self.model, MODEL_REGISTRY.get('CNN_Li2020')) or isinstance(self.model, MODEL_REGISTRY.get('SimpleC3D'))):
        for layer_name in layer_names:
            layer = utils.get_layer(self.model, layer_name)
            if isinstance(layer, torch.nn.Conv3d):
                fig = utils.draw_conv2d_weight(layer.weight)
                image_tensor = utils.fig2rgb(fig)
                save_name = (tag + f'/unscaled/{layer_name}')
                self.logger.experiment.add_image(save_name, image_tensor, step)
                save_name += f'/{step}.png'
                mlflow.log_figure(fig, save_name)
                fig = utils.draw_conv2d_weight(layer.weight, vmin=(- 0.3), vmax=0.3)
                image_tensor = utils.fig2rgb(fig)
                save_name = (tag + f'/uniform_scaled/{layer_name}')
                self.logger.experiment.add_image(save_name, image_tensor, step)
                save_name += f'/{step}.png'
                mlflow.log_figure(fig, save_name)