파이토치 무작정 시작하기 6 - 모델 변수 프리징

 

 

 

 

 

https://www.youtube.com/watch?v=WCByVWpf4rQ

 

 

# 라이브러리 임포트   import torch import torchvision import torchvision.transforms as transforms   # data preprocessing from torch.utils.data import DataLoader   # mini-batch   import torch.nn as nn   # loss import torch.optim as optim   # optimizer
# 구글 드라이브와 구글 코랩을 연동   from google.colab import drive drive.mount("/content/gdrive")
# 데이터 경로 지정 및 GPU 설정   path = '/content/gdrive/My Drive/cifar10' device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(device)
# traindata, testdata 불러오기   transform = transforms.Compose(     [transforms.ToTensor(),      transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))])   trainset = torchvision.datasets.CIFAR10(root=path+'data', train=True, download=True, transform=transform)   trainloader = torch.utils.data.DataLoader(trainset,batch_size=100, shuffle=True, num_workers=2)   testset = torchvision.datasets.CIFAR10(root=path+'data', train=False, download=True, transform=transform)   testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)
# 데이터 확인하기   import matplotlib.pyplot as plt import numpy as np   def imshow(img):   img = img / 2 + 0.5   npimg = img.numpy()   fig = plt.figure(figsize=(10,5))   plt.imshow(np.transpose(npimg, (1, 2, 0)))   plt.show()   dataiter = iter(trainloader) images, labels = dataiter.next()   imshow(torchvision.utils.make_grid(images[:8]))
# 모델 불러오기 # pretrained=True : Imagenet 데이터로 미리 학습한 모델 가져오기 (False면 그냥 모델 구조만 가져옴) # to(device) : 매우중요!!없으면 cpu로 연산함   model = torchvision.models.resnet50(pretrained=True).to(device)
# 불러온 모델 구조 확인   model
# output을 10개로 바꿔주기   num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 10) model = model.to(device)
# 모델의 name을 뽑아내고,150번째에 오면 param.requires_grad = False를 멈춘다. i = 0 for name, param in model.named_parameters():   i += 1   print(i, name)   param.requires_grad = False   if i == 150 :     break
# 특정 레이어를 프리징 시키기 # layer3.0.으로 쓰면 오류남 model.layer3[0].conv3.requires_grad = False model.layer4.requires_grad = False
# loss함수와 optimizer 정의   criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=1e-3)
# training loss를 기준으로 모델 업데이트   old_loss = 10 PATH = path+'cifar_model.pth'   for epoch in range(5):     running_loss = 0.0   for i, data in enumerate(trainloader, 0):          inputs, labels = data[0].to(device), data[1].to(device)       optimizer.zero_grad()     outputs = model(inputs)     loss = criterion(outputs, labels)     loss.backward()     optimizer.step()       running_loss += loss.item()     cost = running_loss / len(trainloader)   print('[%d] loss : %.3f' %(epoch + 1, cost))   if cost < old_loss:     print("Save the best model")     torch.save(model.state_dict(), PATH)     old_loss = cost   print('Finished Training')