权重生成
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
BATCH_SIZE = 12
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCHS = 10
pipeline = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081, ))])
from torch.utils.data import DataLoader
train_set = datasets.MNIST(root="./mnist_data/", train=True, transform=pipeline, download=True)
test_set = datasets.MNIST(root="./mnist_data/", train=False, transform=pipeline, download=True)
train_loader = DataLoader(dataset=train_set, batch_size=BATCH_SIZE, shuffle=False)
test_loader = DataLoader(dataset=test_set, batch_size=BATCH_SIZE, shuffle=True)
import matplotlib.pyplot as plt
images, labels = next(iter(train_loader))
class Digit2(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 10, 5)
self.conv2 = nn.Conv2d(10, 20, 3)
self.fc1 = nn.Linear(20*10*10, 500)
self.fc2 = nn.Linear(500, 10)
self.relu1 = nn.ReLU()
self.relu2 = nn.ReLU()
self.relu3 = nn.ReLU()
self.maxpool2d1 = nn.MaxPool2d(2, 2)
self.maxpool2d2 = nn.MaxPool2d(2, 2)
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
input_size = x.size(0)
x = self.conv1(x)
x = self.relu1(x)
x = self.maxpool2d1(x)
x = self.conv2(x)
x = self.relu2(x)
x = x.view(input_size, -1)
x = self.fc1(x)
x = self.relu3(x)
x = self.fc2(x)
out = self.softmax(x)
return out
model = Digit2().to(DEVICE)
optimizer = optim.Adam(model.parameters())
def train_model(model, device, train_loader, optimizer, epoch):
model.train()
for batch_index, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.cross_entropy(output, target)
loss.backward()
optimizer.step()
if batch_index % 100 == 0:
print("Train Epoch:{} \t loss:{:.6f}".format(epoch, loss.item()))
def test_model(model, device, test_loader):
model.eval()
correct = 0.0
test_loss = 0.0
with torch.no_grad():
for batch_index, (data, target) in enumerate(test_loader):
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.cross_entropy(output, target).item()
pred = output.argmax(dim=1)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print("Test--Average loss:{:.4f}, Accuracy: {:.3f}\n".format(test_loss, 100.0*correct/len(test_loader.dataset)))
for epoch in range(1, EPOCHS+1):
train_model(model, DEVICE, train_loader, optimizer, epoch)
test_model(model, DEVICE, test_loader)
torch.save(model.state_dict(), "params_w.pth")
权重加载
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
BATCH_SIZE = 12
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCHS = 10
pipeline = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081, ))])
from torch.utils.data import DataLoader
train_set = datasets.MNIST(root="./mnist_data/", train=True, transform=pipeline, download=True)
test_set = datasets.MNIST(root="./mnist_data/", train=False, transform=pipeline, download=True)
train_loader = DataLoader(dataset=train_set, batch_size=BATCH_SIZE, shuffle=False)
test_loader = DataLoader(dataset=test_set, batch_size=BATCH_SIZE, shuffle=True)
import matplotlib.pyplot as plt
images, labels = next(iter(train_loader))
class Digit2(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 10, 5)
self.conv2 = nn.Conv2d(10, 20, 3)
self.fc1 = nn.Linear(20*10*10, 500)
self.fc2 = nn.Linear(500, 10)
self.relu1 = nn.ReLU()
self.relu2 = nn.ReLU()
self.relu3 = nn.ReLU()
self.maxpool2d1 = nn.MaxPool2d(2, 2)
self.maxpool2d2 = nn.MaxPool2d(2, 2)
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
input_size = x.size(0)
x = self.conv1(x)
x = self.relu1(x)
x = self.maxpool2d1(x)
x = self.conv2(x)
x = self.relu2(x)
x = x.view(input_size, -1)
x = self.fc1(x)
x = self.relu3(x)
x = self.fc2(x)
out = self.softmax(x)
class Digit3(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 10, 5)
self.conv2 = nn.Conv2d(10, 20, 3)
self.fc1 = nn.Linear(20*10*10, 500)
self.fc3 = nn.Linear(500, 50)
self.fc4 = nn.Linear(50, 10)
self.relu1 = nn.ReLU()
self.relu2 = nn.ReLU()
self.relu3 = nn.ReLU()
self.relu4 = nn.ReLU()
self.maxpool2d1 = nn.MaxPool2d(2, 2)
self.maxpool2d2 = nn.MaxPool2d(2, 2)
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
input_size = x.size(0)
x = self.conv1(x)
x = self.relu1(x)
x = self.maxpool2d1(x)
x = self.conv2(x)
x = self.relu2(x)
x = x.view(input_size, -1)
x = self.fc1(x)
x = self.relu3(x)
x = self.fc3(x)
x = self.relu4(x)
x = self.fc4(x)
out = self.softmax(x)
return out
digit2 = Digit2().to(DEVICE)
digit3 = Digit3().to(DEVICE)
optimizer = optim.Adam(digit3.parameters())
digit2.load_state_dict(torch.load("params_w.pth"))
pretrained_dict = digit2.state_dict()
model_dict = digit3.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
digit3.load_state_dict(model_dict)
torch.save(digit3.state_dict(), "params_new.pth")
digit3.load_state_dict(torch.load("params_new.pth"))
model_dict1 = digit3.state_dict()
print(1)