如何在Python中将Alexnet的渐变存储为numpy数组(在每次迭代中)？

pip install numpy torch torchvision

import torch
import torchvision.models as models
import numpy as np

model = models.alexnet(pretrained=True)

def get_gradients(model):
    gradients = []
    for param in model.parameters():
        if param.grad is not None:
            gradients.append(param.grad.data.numpy())
    return gradients

# 假设有一个输入张量input_tensor和目标张量target_tensor
input_tensor = torch.randn(1, 3, 224, 224)
target_tensor = torch.randn(1, 1000)

# 将输入张量传递给模型并计算输出
output_tensor = model(input_tensor)

# 计算损失
loss = torch.nn.functional.mse_loss(output_tensor, target_tensor)

# 清除之前的渐变
model.zero_grad()

# 计算渐变
loss.backward()

# 获取渐变并存储为numpy数组
gradients = get_gradients(model)
gradients_np = [grad for grad in gradients]