保存网络中的参数到txt

import os
import scipy.io as io

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader

# from pytorchtools import EarlyStopping

from torchvision.datasets import MNIST, CIFAR10, CIFAR100
from torchvision.transforms import Compose, ToTensor, Normalize
from torchvision import transforms


VGG_CONFIGS = {
    # M for MaxPool, Number for channels
    'D': [
        64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M',
        512, 512, 512, 'M'
    ],
    'E': [
        64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M',
        512, 512, 512, 512, 'M'
    ],
}


class VGG_SNIP(nn.Module):
    """
    This is a base class to generate three VGG variants used in SNIP paper:
        1. VGG-C (16 layers)
        2. VGG-D (16 layers)
        3. VGG-like

    Some of the differences:
        * Reduced size of FC layers to 512
        * Adjusted flattening to match CIFAR-10 shapes
        * Replaced dropout layers with BatchNorm
    """

    def __init__(self, config, num_classes=100):
        super().__init__()

        self.features = self.make_layers(VGG_CONFIGS[config], batch_norm=True)

        self.classifier = nn.Sequential(
            nn.Linear(512, 512),  # 512 * 7 * 7 in the original VGG
            nn.ReLU(True),
            nn.BatchNorm1d(512),  # instead of dropout
            nn.Linear(512, 512),
            nn.ReLU(True),
            nn.BatchNorm1d(512),  # instead of dropout
            nn.Linear(512, num_classes),
        )

    @staticmethod
    def make_layers(config, batch_norm=False):  # TODO: BN yes or no?
        layers = []
        in_channels = 3
        for v in config:
            if v == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
                if batch_norm:
                    layers += [
                        conv2d,
                        nn.BatchNorm2d(v),
                        nn.ReLU(inplace=True)
                    ]
                else:
                    layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels = v
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        x = F.log_softmax(x, dim=1)
        return x

net = VGG_SNIP('D').cuda()
net = torch.nn.DataParallel(net).cuda()
checkpoint = torch.load(r'E:\code\practicing\snip-master\VGG16_2.pth.tar')
net.load_state_dict(checkpoint['state_dict'])



for name, parameter in net.named_parameters():
    with open('C:\\Users\\Pong\\Desktop\\for_ly\\{}.txt'.format(name), 'w') as thisfile:
        if (len(parameter.shape))==1:
            for channel in parameter:
                channel = str(channel.cpu().detach().numpy())
                thisfile.write(channel)
                thisfile.write("\n")
        if (len(parameter.shape))==2:
            for channel in parameter:
                channel = channel.cpu().detach().numpy()
                np.savetxt(thisfile, channel, fmt='%f', delimiter = ',')
                thisfile.write("\n")
        elif (len(parameter.shape))==4:
            for channel in parameter:
                for filter in channel:
                    filter = filter.cpu().detach().numpy()
                    filter = np.array(filter)
                    np.savetxt(thisfile, filter, fmt='%f', delimiter = ',')
                    thisfile.write("\n")
    # c = np.loadtxt('/home/Velocitymodel/speedfile/b.txt', delimiter = ',').reshape((2, 2, 2))  恢复


print(":)")