# -*- coding: utf-8 -*-
"""
Auther : Haitao Zeng
date : 2019.2.14
Function: finetune the pre-trained model on MIT67 & SUN397
"""
from __future__ import print_function, division, absolute_import
from PIL import Image
import torch
import torchvision
import torch.nn as nn
from torch.autograd import Variable
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
from torchvision import datasets, models, transforms
import os
import cv2
import time
import copy
import torch.utils.data as data
from Rsenet50 import Resnet
import torch.utils.model_zoo as model_zoo
def cv2_imageloader(path):
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
img = cv2.imread(path)
img = cv2.resize(img, (224, 224))
im_arr = np.float32(img)
im_arr = np.ascontiguousarray(im_arr[..., ::-1])
im_arr = im_arr.transpose(2, 0, 1)# Convert Img from BGR to RGB
for channel, _ in enumerate(im_arr):
# Normalization
im_arr[channel] /= 255
im_arr[channel] -= mean[channel]
im_arr[channel] /= std[channel]
# Convert to float tensor
im_as_ten = torch.from_numpy(im_arr).float()
# Convert to Pytorch variable
im_as_var = Variable(im_as_ten, requires_grad=True)
return im_as_var
def default_loader(path):
return cv2_imageloader(path)
class CustomImageLoader(data.Dataset):
##自定义类型数据输入
def __init__(self, img_path, txt_path, dataset = '', data_transforms=None, loader = default_loader):
im_list = []
im_labels = []
with open(txt_path, 'r') as files:
for line in files:
items = line.split()
if items[0][0] == '/':
# fnewname = '_'.join(items[0][1:-4].split('/'))
imname = line.split()[0][1:]
else:
# fnewname = '_'.join(items[0][:-4].split('/'))
imname = line.split()[0]
im_list.append(os.path.join(img_path, imname))
im_labels.append(int(items[1]))
self.imgs = im_list
self.labels = im_labels
# self.data_tranforms = data_transforms
self.loader = loader
self.dataset = dataset
def __len__(self):
return len(self.imgs)
def __getitem__(self, item):
# print(item)
img_name = self.imgs[item]
label = self.labels[item]
img = self.loader(img_name)
return img, label
NUM_EPOCH = 20
batch_size = 32
device = torch.device('cuda:0')
#MIT67 INPUT
image_dir = '/media/haitaizeng/00038FCE000387A5/cgw/Datasets/MIT67/Images'
image_datasets = {x : CustomImageLoader(image_dir, txt_path=('/home/haitaizeng/stanforf/alex_mit/data_image/'+x+'Images.label'),
data_transforms=data_tranforms,
dataset=x) for x in ['Train', 'Test']
}
dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=batch_size,
shuffle=True) for x in ['Train', 'Test']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['Train', 'Test']}
def train_model(model, crtiation, optimizer,schedular, num_epochs=NUM_EPOCH):
begin_time = time.time()
best_weights = copy.deepcopy(model.state_dict())#copy the weights from the model
best_acc = 0.0
for epoch in range(num_epochs):
print("-*-" * 20)
for phase in ['Train', 'Test']:
if phase=='Train':
schedular.step()
model.train()
else:
model.eval()
running_loss = 0.0
running_acc = 0.0
for images, labels in dataloders[phase]:
images.to(device)
labels.to(device)
optimizer.zero_grad()
with torch.set_grad_enabled(phase=='Train'):
opt = model(images.cuda())
# opt = model(images)
_,pred = torch.max(opt,1)
labels = labels.cuda()
loss = crtiation(opt, labels)
if phase=='Train':
loss.backward()
optimizer.step()
running_loss += loss.item()*images.size(0)
running_acc += torch.sum(pred==labels)
epoch_loss = running_loss/dataset_sizes[phase]
epoch_acc = running_acc.double()/dataset_sizes[phase]
print('epoch={}, Phase={}, Loss={:.4f}, ACC:{:.4f}'.format(epoch, phase, epoch_loss, epoch_acc))
if phase == 'Test' and epoch_acc>best_acc and epoch_acc > 0.8:
# Upgrade the weights
best_acc=epoch_acc
best_weights = copy.deepcopy(model.state_dict())
### Save model_dict
path ='/home/haitaizeng/stanforf/ctumb_zht/TFP/pytorch/save_model/MIT67'
PATH = os.path.join(path,"Places"+ str("{:.4f}".format(epoch_acc))+".pth")
print(PATH)
torch.save(model.state_dict(), PATH)
time_elapes = time.time() - begin_time
print('Training Complete in{:.0f}m {:0f}s'.format(
time_elapes // 60, time_elapes % 60
))
print('Best Val ACC: {:}'.format(best_acc))
model.load_state_dict(best_weights)
return models
if __name__ == '__main__':
# NUMCLASS = 397
NUMCLASS=67
pthpath = '/home/haitaizeng/cgw/for_zht/CVPR19/resnet50_places365.pth.tar'
# model_ft = models.__ dict__['resnet50'](num_classes=365) ## 加载的是pytorch库中预先写好的Resnet50的网络结构
model_ft = Resnet([3, 4, 6, 3], 365)##这是自行编写的Resnet50，用于后面的特征提取的操作
##load pre-trained model for funetuning
ckpt = torch.load(pthpath, map_location=lambda storage, loc: storage)
state_dict = {str.replace(k, 'module.', ''): v for k, v in ckpt['state_dict'].items()}
model_ft.load_state_dict(state_dict)
## change the last fully connedected layer
num_fits = model_ft.fc.in_features
### 此处功能是固定前面的层不动，只是Finetune最后一层全连接层。#固定前面的层不进行反向传播计算
params_to_update = model_ft.parameters()
# for parm in model_ft.parameters():
# parm.requires_grad=False
model_ft.fc = nn.Linear(num_fits, NUMCLASS)
feature_extract = False
# 检查需要Finetune的layer名称
if feature_extract is True:
params_to_update=[]
for name, param in model_ft.named_parameters():
if param.requires_grad==True:
params_to_update.append(param)
print('\t', name)
else:
for name, param in model_ft.named_parameters():
if param.requires_grad==True:
print('upgrade layer:', name)
# else:
# print("No upgrade param", name)
## 训练这个模型在cuda上 the model on the cuda---Nvidia -1080ti
model_ft = model_ft.to(device)
model_ft.cuda()
## cacluate the cross entropu loss
criterion = nn.CrossEntropyLoss()
## using SGD to optimize,此处是finetune了所有的层的结果
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
## learning rate decay ,when epoch ==10
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)
model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)
