本來是只用Tenorflow的，但是因為TF有些Numpy特性並不支持，比如對數組使用列表進行切片，所以只能轉戰Pytorch了（pytorch是支持的）。還好Pytorch比較容易上手，幾乎完美復制了Numpy的特性（但還有一些特性不支持），怪不得熱度上升得這么快。
模型定義
和TF很像，Pytorch也通過繼承父類來搭建自定義模型，同樣也是實現兩個方法。在TF中是__init__()和call()，在Pytorch中則是__init__()和forward()。功能類似，都分別是初始化模型內部結構和進行推理。其它功能比如計算loss和訓練函數，你也可以繼承在里面，當然這是可選的。下面搭建一個判別MNIST手寫字的Demo，首先給出模型代碼：
import numpy as np import matplotlib.pyplot as plt import torch from torch import nn,optim from torchsummary import summary from keras.datasets import mnist from keras.utils import to_categorical device = torch.device('cuda') #——————1—————— class ModelTest(nn.Module): def __init__(self,device): super().__init__() self.layer1 = nn.Sequential(nn.Flatten(),nn.Linear(28*28,512),nn.ReLU())#——————2—————— self.layer2 = nn.Sequential(nn.Linear(512,512),nn.ReLU()) self.layer3 = nn.Sequential(nn.Linear(512,512),nn.ReLU()) self.layer4 = nn.Sequential(nn.Linear(512,10),nn.Softmax()) self.to(device) #——————3—————— self.opt = optim.SGD(self.parameters(),lr=0.01)#——————4—————— def forward(self,inputs): #——————5—————— x = self.layer1(inputs) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) return x def get_loss(self,true_labels,predicts): loss = -true_labels * torch.log(predicts) #——————6—————— loss = torch.mean(loss) return loss def train(self,imgs,labels): predicts = model(imgs) loss = self.get_loss(labels,predicts) self.opt.zero_grad()#——————7—————— loss.backward()#——————8—————— self.opt.step()#——————9—————— model = ModelTest(device) summary(model,(1,28,28),3,device='cuda') #——————10——————
#1：獲取設備，以方便后面的模型與變量進行內存遷移，設備名只有兩種：'cuda'和'cpu'。通常是在你有GPU的情況下需要這樣顯式進行設備的設置，從而在需要時，你可以將變量從主存遷移到顯存中。如果沒有GPU，不獲取也沒事，pytorch會默認將參數都保存在主存中。
#2：模型中層的定義，可以使用Sequential將想要統一管理的層集中表示為一層。
#3：在初始化中將模型參數遷移到GPU顯存中，加速運算，當然你也可以在需要時在外部執行model.to(device)進行遷移。
#4：定義模型的優化器，和TF不同，pytorch需要在定義時就將需要梯度下降的參數傳入，也就是其中的self.parameters()，表示當前模型的所有參數。實際上你不用擔心定義優化器和模型參數的順序問題，因為self.parameters()的輸出並不是模型參數的實例，而是整個模型參數對象的指針，所以即使你在定義優化器之后又定義了一個層，它依然能優化到。當然優化器你也可以在外部定義，傳入model.parameters()即可。這里定義了一個隨機梯度下降。
#5：模型的前向傳播，和TF的call()類似，定義好model()所執行的就是這個函數。
#6：我將獲取loss的函數集成在了模型中，這里計算的是真實標簽和預測標簽之間的交叉熵。
#7/8/9：在TF中，參數梯度是保存在梯度帶中的，而在pytorch中，參數梯度是各自集成在對應的參數中的，可以使用tensor.grad來查看。每次對loss執行backward()，pytorch都會將參與loss計算的所有可訓練參數關於loss的梯度疊加進去（直接相加）。所以如果我們沒有疊加梯度的意願的話，那就要在backward()之前先把之前的梯度刪除。又因為我們前面已經把待訓練的參數都傳入了優化器，所以，對優化器使用zero_grad()，就能把所有待訓練參數中已存在的梯度都清零。那么梯度疊加什么時候用到呢？比如批量梯度下降，當內存不夠直接計算整個批量的梯度時，我們只能將批量分成一部分一部分來計算，每算一個部分得到loss就backward()一次，從而得到整個批量的梯度。梯度計算好后，再執行優化器的step()，優化器根據可訓練參數的梯度對其執行一步優化。
#10：使用torchsummary函數顯示模型結構。奇怪為什么不把這個繼承在torch里面，要重新安裝一個torchsummary庫。
訓練及可視化
接下來使用模型進行訓練，因為pytorch自帶的MNIST數據集並不好用，所以我使用的是Keras自帶的，定義了一個獲取數據的生成器。下面是完整的訓練及繪圖代碼（50次迭代記錄一次准確率）：
import numpy as np import matplotlib.pyplot as plt import torch from torch import nn,optim from torchsummary import summary from keras.datasets import mnist from keras.utils import to_categorical device = torch.device('cuda') #——————1—————— class ModelTest(nn.Module): def __init__(self,device): super().__init__() self.layer1 = nn.Sequential(nn.Flatten(),nn.Linear(28*28,512),nn.ReLU())#——————2—————— self.layer2 = nn.Sequential(nn.Linear(512,512),nn.ReLU()) self.layer3 = nn.Sequential(nn.Linear(512,512),nn.ReLU()) self.layer4 = nn.Sequential(nn.Linear(512,10),nn.Softmax()) self.to(device) #——————3—————— self.opt = optim.SGD(self.parameters(),lr=0.01)#——————4—————— def forward(self,inputs): #——————5—————— x = self.layer1(inputs) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) return x def get_loss(self,true_labels,predicts): loss = -true_labels * torch.log(predicts) #——————6—————— loss = torch.mean(loss) return loss def train(self,imgs,labels): predicts = model(imgs) loss = self.get_loss(labels,predicts) self.opt.zero_grad()#——————7—————— loss.backward()#——————8—————— self.opt.step()#——————9—————— def get_data(device,is_train = True, batch = 1024, num = 10000): train_data,test_data = mnist.load_data() if is_train: imgs,labels = train_data else: imgs,labels = test_data imgs = (imgs/255*2-1)[:,np.newaxis,...] labels = to_categorical(labels,10) imgs = torch.tensor(imgs,dtype=torch.float32).to(device) labels = torch.tensor(labels,dtype=torch.float32).to(device) i = 0 while(True): i += batch if i > num: i = batch yield imgs[i-batch:i],labels[i-batch:i] train_dg = get_data(device, True,batch=4096,num=60000) test_dg = get_data(device, False,batch=5000,num=10000) model = ModelTest(device) summary(model,(1,28,28),11,device='cuda') ACCs = [] import time start = time.time() for j in range(20000): #訓練 imgs,labels = next(train_dg) model.train(imgs,labels) #驗證 img,label = next(test_dg) predicts = model(img) acc = 1 - torch.count_nonzero(torch.argmax(predicts,axis=1) - torch.argmax(label,axis=1))/label.shape[0] if j % 50 == 0: t = time.time() - start start = time.time() ACCs.append(acc.cpu().numpy()) print(j,t,'ACC: ',acc) #繪圖 x = np.linspace(0,len(ACCs),len(ACCs)) plt.plot(x,ACCs)
准確率變化圖如下：
其它使用技巧
tensor與array
需要注意的是，pytorch的tensor基於numpy的array，它們是共享內存的。也就是說，如果你把tensor直接插入一個列表，當你修改這個tensor時，列表中的這個tensor也會被修改；更容易被忽略的是，即使你用tensor.detach.numpy()，先將tensor轉換為array類型，再插入列表，當你修改原本的tensor時，列表中的這個array也依然會被修改。所以如果我們只是想保存tensor的值而不是整個對象，就要使用np.array(tensor)將tensor的值復制出來。
自定義層
在TF中，自定義模型通常繼承keras的Model，而自定義層則是繼承layers.Layer，繼承不同的父類通常會造成初學者的困擾。而在pytorch中，自定義層與自定義模型一樣，都是繼承nn.Module。Pytorch將層與模型都看成了模塊，這很容易理解。的確，層與模型之間本來也沒有什么明確的界限。並且定義方式與上面定義模型的方式一樣，也是實現兩個函數即可。代碼示例如下：
import torch from torch import nn class ParaDeconv(nn.Module):#——————1—————— def __init__(self,in_n,out_n): super().__init__() self.w = nn.Parameter(torch.normal(0,0.01,size = [in_n,out_n]),requires_grad=True) self.b = nn.Parameter(torch.normal(0,0.01,size = [out_n]),requires_grad=True) def forward(self,inputs): x = torch.matmul(inputs,self.w) x = x + self.b return x layer = ParaDeconv(2,3) y = layer(torch.ones(100,2))#——————2—————— loss = torch.sum(y)#——————3—————— loss.backward()#——————4—————— for i in layer.parameters():#——————5—————— print(i.grad)#——————6——————
#1：自定義一個全連接層。層中可訓練參數的定義是使用nn.Parameter，如果直接使用torch.tensor是無法在#5中遍歷到的。
#2/3/4：輸入並計算loss，然后反向傳播計算參數梯度。
#5/6：輸出完成反向傳播后層參數的梯度。
以上定義的層可以和pytorch自帶的層一樣直接插入模型中使用。
保存/加載
保存/加載模型
有兩種方式，一種是保存模型的參數：
torch.save(model.state_dict(), PATH) #保存 model.load_state_dict(torch.load(PATH),strict=True) #加載
這種加載方式需要先定義模型，然后再加載參數。如果你定義的模型參數名與保存的參數對不上，就會出錯。但如果把strict修改成False，不嚴格匹配，它就會只匹配對應上的鍵值，不會因多出或缺少的參數而報錯。
另一種是直接保存模型：
torch.save(model, PATH) #保存 model = torch.load(PATH) #加載
這種方式看似方便，實際上更容易出錯。因為python不能保存整個模型的類，所以它只能保存定義類的代碼文件位置，以在加載時獲取類的結構。如果你改變了定義類的代碼位置，就有可能因為找不到類而出錯。
保存訓練點
當你要保存某個訓練階段的狀態，比如包含優化器參數、模型參數、訓練迭代次數等，可以進行如下操作：
#保存訓練點 torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss }, PATH) #加載訓練點 model = TheModelClass(*args, **kwargs) optimizer = TheOptimizerClass(*args, **kwargs) checkpoint = torch.load(PATH) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) epoch = checkpoint['epoch'] loss = checkpoint['loss']
和保存模型一樣，也是使用torch.save()。它很靈活，可以保存字典，因此讀取的時候也按照字典索引讀取即可。當然要注意，並不是任何類型都能保存的，這里保存的四個類型分別是：
1. int
2. collections.OrderedDict
3. collections.OrderedDict
4. list
修改模型參數
Pytorch沒有提供額外的方式讓我們修改模型參數，我們可以使用上面加載模型參數的方式來修改參數。對於某個參數，我們只要把鍵值和對應要修改的值放在字典中傳入load_state_dict即可。如果沒傳入所有的參數，記得把strict設為False。示例如下：
model.load_state_dict({'weight':torch.tensor([0.])},strict=False) #修改模型參數
參數名，也就是鍵值，和對應的參數shape可以通過model.state_dict()查看。
添加Grad Penalty (GP)
WGAN-GP的grad penalty需要用到二次求導，直接無腦使用backward()方法是不行的。梯度的計算有兩種方法，一種是常見的backward()，另一種是autograd.grad()。
下面做一個簡單的實驗來記錄如何使用這兩者計算GP。代碼如下：
import torch from torch import nn, optim from time import time from torch import autograd torch.manual_seed(1) #————0———— model = nn.Linear(3, 2) #————0———— opt = optim.SGD(model.parameters(), 0.1) #————0———— start_t = time() for i in range(10): x = torch.ones([1, 3], requires_grad=True) #————1———— y = torch.sum(model(x)) #————1———— # *********************** y.backward(retain_graph = False, create_graph = True) #————2———— g = x.grad #————2———— # g = autograd.grad(y, x, retain_graph = False, create_graph = True)[0] # *********************** print("Input gradient: \n", g.detach().numpy())#————3———— for i in model.state_dict(): #————3———— print(i+":\n", model.state_dict()[i].numpy()) #————3———— gp = torch.sum(g**2) #————4———— print("GP: \n", gp.detach().numpy()) print() opt.zero_grad() #————5———— gp.backward() #————5———— opt.step() #————5———— print("Cost time: \n", time() - start_t)
注釋：
0、設置隨機種子、創建只包含一個全連接層的模型($R^3\to R^2$)、創建優化器。
1、定義可獲取梯度的輸入，並通過模型求得輸出。
2、使用backward()或autograd.grad()獲取輸入關於輸出的梯度。需要注意的是，create_graph = True 表示在反向傳播計算梯度時記錄計算圖。因為梯度通常直接用來梯度下降，無需關於梯度的計算圖，所以默認為False。而我們要用這個梯度再一次反向傳播來計算Grad Penalty，所以需要設置為True。另外，retain_graph設置為False，表示計算完梯度后，關於模型輸出的前向傳播計算圖就釋放掉。因為我們需要的是輸入關於輸出的梯度的計算圖，因此在獲取梯度計算圖之后，輸入到輸出的計算圖就可以釋放掉了。很多博客都設置為True，這是在浪費資源。
3、打印輸入關於輸出的梯度，並打印模型的權重值。可以驗證，打印出的梯度向量值等於權重行向量之和。
4、計算GP。為了便於理解，這里直接用二范數的平方來代替。
5、使用GP更新模型權重。
輸出結果如下：
可以看出權重兩個行向量之和越來越接近0。這和數學計算的預期結果是一致的。
以上列舉的兩個方法，是有計算量上的差異的。backward()會計算所有與輸出相關，且可獲取梯度的參數的梯度；而autograd.grad()則只會計算輸入與輸出之間，與輸入和輸出都相關的梯度，且函數返回的只有輸入關於輸出的梯度。也就是說，在以上實驗中，autograd.grad()不會計算全連接層中bias關於輸出的梯度（因為它並沒有影響到輸入關於輸出的梯度）。所以，使用autograd.grad()計算GP會更快。
我們可以使用以上代碼進行實驗，兩種方法在10000次迭代中分別用時3.5s和3.2s，autograd.grad()快快了0.3秒。而這里的bias規模只有2，當規模大起來，節省的時間就很可觀了。