一個完整的Pytorch深度學習專案程式碼
Network
建立一個Network類,繼承torch.nn.Module,在建構函式中用初始化成員變數為具體的網路層,在forward函式中使用成員變數搭建網路架構,模型的使用過程中pytorch會自動呼叫forword進行引數的前向傳播,構建計算圖。以下拿一個簡單的CNN影象分類模型舉例:
class Network(nn.Module):
def __init__(self):
super(Network, self).__init__()
# 灰度影象的channels=1即in_channels=1 輸出為10個類別即out_features=10
# parameter(形參)=argument(實參) 卷積核即卷積濾波器 out_channels=6即6個卷積核 輸出6個feature-maps(特徵對映)
# 權重shape 6*1*5*5
self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
self.bn1 = nn.BatchNorm2d(6) # 二維批歸一化 輸入size=6
# 權重shape 12*1*5*5
self.conv2 = nn .Conv2d(in_channels=6, out_channels=12, kernel_size=5)
# 全連線層:fc or dense or linear out_features即特徵(一階張量)
# 權重shape 120*192
self.fc1 = nn.Linear(in_features=12*4*4, out_features=120)
self.bn2 = nn.BatchNorm1d(120) # 一維批歸一化 輸入size=120
# 權重shape 60*120
self .fc2 = nn.Linear(in_features=120, out_features=60)
# 權重shape 10*60
self.out = nn.Linear(in_features=60, out_features=10)
def forward(self, t):
# (1) input layer
t = t
# (2) hidden conv layer
t = F.relu(self.conv1(t)) # (28-5+0)/1+1=24 輸入為b(batch_size)*1*28*28 輸出為b*6*24*24 relu後shape不變
t = F.max_pool2d(t, kernel_size=2, stride=2) # (24-2+0)/2+1=12 輸出為b*6*12*12
t = self.bn1(t)
# (3) hidden conv layer
t = F.relu(self.conv2(t)) # (12-5+0)/1+1=8 輸出為b*12*8*8 relu後shape不變
t = F.max_pool2d(t, kernel_size=2, stride=2) # (8-2+0)/2+1=4 輸出為b*12*4*4
# (4) hidden linear layer
t = F.relu(self.fc1(t.reshape(-1, 12*4*4))) # t.reshape後為b*192 全連線層後輸出為b*120 relu後shape不變
t = self.bn2(t)
# (5) hidden linear layer
t = F.relu(self.fc2(t)) # 全連線層後輸出為b*60 relu後shape不變
# (6) output layer
t = self.out(t) # 全連線層後輸出為b*10 relu後shape不變
return t
Transofrms
資料處理可以直接使用torchvision.transforms下的處理函式,包括均值,隨機旋轉,隨機裁剪等等,也可以自己實現一些pytorch中沒有實現的處理函式,下面拿一個分割網路的處理函式舉例,可支援同時對傳入的Image和GroundTruth進行處理,使用時直接按照順序構造ProcessImgAndGt即可。
class ProcessImgAndGt(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, img, label):
for t in self.transforms:
img, label = t(img, label)
return img, label
class Resize(object):
def __init__(self, height, width):
self.height = height
self.width = width
def __call__(self, img, label):
img = img.resize((self.width, self.height), Image.BILINEAR)
label = label.resize((self.width, self.height), Image.NEAREST)
return img, label
class Normalize(object):
def __init__(self, mean, std):
self.mean, self.std = mean, std
def __call__(self, img, label):
for i in range(3):
img[:, :, i] -= float(self.mean[i])
for i in range(3):
img[:, :, i] /= float(self.std[i])
return img, label
class ToTensor(object):
def __init__(self):
self.to_tensor = torchvision.transforms.ToTensor()
def __call__(self, img, label):
img, label = self.to_tensor(img), self.to_tensor(label).long()
return img, label
transforms = ProcessImgAndGt([
Resize(512, 512),
Normalize([0.5, 0.5, 0.5], [0.1, 0.1, 0.1]),
ToTensor()
])
Dataset
建立一個數據集類,繼承torch.utils.data.Dataset,只需實現__init__建構函式,__getitem__迭代器遍歷函式以及__len__函式。
- 在__init__函式中讀取傳入的資料集路徑下的指定資料檔案,還是拿一個分割網路的dataset流程舉例,其他分類分類模型可以直接將GroundTruth替換為對應label即可,將拼接處理好的圖片檔案路徑和GroundTruth檔案路徑作為元組存入一個為列表的成員變數file_list中;
- 在__getitem__中根據傳入的索引從file_list取對應的元素,並且通過Transforms進行處理;
- 在__len__中返回len(self.file_list)即可。
class MyDataset(torch.utils.data.Dataset):
def __init__(self, dataset_path, transforms):
super(TrainDataset, self).__init()
self.dataset_path = dataset_path
self.transforms = transforms
# 根據具體的業務邏輯讀取全部資料路徑作為載入資料的索引
for dir in os.listdir(dataset_path):
image_dir = os.path.join(dataset_path, dir)
gt_path = image_dir + '/GT/'
img_path = image_dir + '/Frame/'
img_list = []
for name in os.listdir(img_path):
if name.endswith('.png'):
img_list.append(name)
self.file_list.extend([(img_path + name, gt_path + name) for name in img_list])
def __getitem__(self, idx):
img_path, label_path = self.file_list[idx]
img = Image.open(img_path).convert('RGB')
label = Image.open(label_path).convert('L')
img, label = self.transforms(img, label)
return img, label
def __len__(self):
return len(self.file_list)
Optimizer
選擇優化器進行模型引數更新,要建立優化器必須給它一個可進行迭代優化的包含了全部引數的列表 然後可以指定針對這些引數的學習率(learning_rate),權重衰減(weight_decay),momentum等,
optimizer = optim.Adam(model.parameters(), lr = 0.0001)
或者是可以指定針對哪些引數執行不一樣的優化策略,根據不同層的name對不同層使用不同的優化策略。列表中的每一項都可以是一個dict,dict中params對應當前項的引數列表,可以對當前項指定學習率或者是衰減策略。對base_params使用的1e-4的學習率,對finetune_params使用1e-3的學習率,對兩者一起使用1e-4的權重衰減
base_params = [params for name, params in model.named_parameters() if ("xxx" in name)]
finetune_params = [params for name, params in model.named_parameters() if ("yyy" in name)]
optimizer = optim.Adam([
{"params": base_params},
{"params": finetune_params, "lr": 1e-3}
], lr=1e-4, weight_decay=1e-4);
Run
基礎元件都寫好了,剩下的就是組成一個完整的模型結構。
- 例項化模型物件,並將其載入到GPU中
- 根據需要構建資料預處理物件,傳入資料集物件中進行讀取資料時的資料處理
- 構建訓練和測試的資料集物件,並將其傳入torch.utils.data.DataLoader,指定batch_size(訓練或測試是每次讀取多少條資料)、shuffle(讀取資料時是否打亂)、num_workers(開啟多少執行緒進行資料載入,為0時(不推薦)用主執行緒在訓練模型時進行資料載入)等引數
- 使用torch.optim.Adam構建優化器物件,這裡根據不同層的name對不同層使用不同的優化策略
- 訓練20個epoch,並且每5個epoch在測試集上跑一遍,這裡只計算了損失,對於其他評價指標直接計算即可
- 根據條件對指定epoch的模型進行儲存
- optimizer.zero_grad() # pytorch會積累梯度,在優化每個batch的權重的梯度之前將之前計算出的每個權重的梯度置0
- loss.backward() # 在最後一個張量上呼叫反向傳播方法,在計算圖中計算權重的梯度
- optimizer.step() # 使用預先設定的學習率等引數根據當前梯度對權重進行更
- model.train() # 保證BN層能夠繼續計算資料的均值和方差並進行更新,保證dropout層會按照設定的引數設定保留啟用單元的概率(保留概率=p)
- model.eval() # BN層會停止計算均值和方差,直接使用訓練時的引數,dropout層利用了訓練好的全部網路連線,不隨機捨棄啟用單元
model = Network().cuda()
# 構建資料預處理
transforms = ProcessImgAndGt([
Resize(512, 512),
Normalize([0.5, 0.5, 0.5], [0.1, 0.1, 0.1]),
ToTensor()
])
# 構建Dataset
train_dataset = MyDataset(train_dataset_path, transforms)
# DataLoader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=12,
shuffle=True,
num_workers=4,
pin_memory=False)
# TestDataset
test_dataset = MyDataset(test_dataset_path, transforms)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=4,
shuffle=True,
num_workers=2,
pin_memory=False)
# optimizer需要傳入全部需要更新的引數名稱,這裡是對不用的引數執行不同的更新策略
base_params = [params for name, params in model.named_parameters() if ("xxx" in name)]
finetune_params = [params for name, params in model.named_parameters() if ("yyy" in name)]
optimizer = torch.optim.Adam([
{"params": base_params, "lr": 1e-3, ...},
{"params": finetune_params, "lr": 1e-4, ...}
])
for epoch in range(20):
model.train()
epoch_loss = 0
for batch in trian_loader:
images. gts = batch[0].cuda(), batch[1].cuda()
preds = model(iamges)
loss = F.cross_entropy(preds, gts)
optimizer.zero_grad() # pytorch會積累梯度,在優化每個batch的權重的梯度之前將之前計算出的每個權重的梯度置0
loss.backward() # 在最後一個張量上呼叫反向傳播方法,在計算圖中計算權重的梯度
optimizer.step() # 使用預先設定的學習率等引數根據當前梯度對權重進行更新
epoch_loss += loss * trian_loader.batch_size
# 計算其他標準
loss = epoch_loss / len(train_loader.dataset)
# .......
# 每隔幾個epoch在測試集上跑一下
if epoch % 5 == 0:
model.eval()
test_epoch_loss = 0
for test_batch in test_loader:
test_images. test_gts = test_batch[0].cuda(), test_batch[1].cuda()
test_preds = model(test_iamges)
loss = F.cross_entropy(test_preds, test_gts)
test_epoch_loss += loss * test_loader.batch_size
# 計算其他標準
test_loss = test_epoch_loss / (len(test_loader.dataset))
# .......
# 根據條件對指定epoch的模型進行儲存 將模型序列化到磁碟的pickle包
if 精度最高:
torch.save(model.stat_dict(), f'{model_path}_{time_index}.pth')
Test
實際使用時需要將訓練好的模型上在輸入資料上執行,這裡以測試集的資料為例,實際情況下只需要初始化模型之後將視訊流中的影象幀作為模型的輸入即可。
torch.no_grad()
- 停止autograd模組的工作,不計算和儲存梯度,一般在用訓練好的模型跑測試集時使用,因為測試集時不需要計算梯度更不會更新梯度。使用後可以加速計算時間,節約gpu的視訊記憶體
test_dataset = MyDataset(test_dataset_path, transforms)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=1,
shuffle=False,
num_workers=2)
model = Network().cuda()
# 對磁碟上的pickle檔案進行解包 將gpu訓練的模型載入到cpu上
model.load_stat_dict(torch.load(model_path, map_location=torch.device('cpu')));
mocel.eval()
with torch.no_grad():
for batch in test_loader:
test_images. test_gts = test_batch[0].cuda(), test_batch[1].cuda()
test_preds = model(test_iamges)
# 儲存模型輸出的圖片
參考: https://www.zhihu.com/question/406133826/answer/2389936678