如何從TensorFlow轉入PyTorch
當我第一次嘗試學習 PyTorch 時,沒幾天就放棄了。和 TensorFlow 相比,我很難弄清 PyTorch 的核心要領。但是隨后不久,PyTorch 發布了一個新版本,我決定重新來過。在第二次的學習中,我開始了解這個框架的易用性。在本文中,我會簡要解釋 PyTorch 的核心概念,為你轉入這個框架提供一些必要的動力。其中包含了一些基礎概念,以及先進的功能如學習速率調整、自定義層等等。
PyTorch 的易用性如何?Andrej Karpathy 是這樣評價的
資源
- 首先要知道的是:PyTorch 的主目錄和教程是分開的。而且因為開發和版本更新的速度過快,有時候兩者之間并不匹配。所以你需要不時查看源代碼:http://pytorch.org/tutorials/。
- 當然,目前網絡上已有了一些 PyTorch 論壇,你可以在其中詢問相關的問題,并很快得到回復:https://discuss.pytorch.org/。
把 PyTorch 當做 NumPy 用
讓我們先看看 PyTorch 本身,其主要構件是張量——這和 NumPy 看起來差不多。這種性質使得 PyTorch 可支持大量相同的 API,所以有時候你可以把它用作是 NumPy 的替代品。PyTorch 的開發者們這么做的原因是希望這種框架可以完全獲得 GPU 加速帶來的便利,以便你可以快速進行數據預處理,或其他任何機器學習任務。將張量從 NumPy 轉換至 PyTorch 非常容易,反之亦然。讓我們看看如下代碼:
- importtorch
- importnumpy asnp
- numpy_tensor =np.random.randn(10,20)
- # convert numpy array to pytorch array
- pytorch_tensor =torch.Tensor(numpy_tensor)
- # or another way
- pytorch_tensor =torch.from_numpy(numpy_tensor)
- # convert torch tensor to numpy representation
- pytorch_tensor.numpy()
- # if we want to use tensor on GPU provide another type
- dtype =torch.cuda.FloatTensor
- gpu_tensor =torch.randn(10,20).type(dtype)
- # or just call `cuda()` method
- gpu_tensor =pytorch_tensor.cuda()
- # call back to the CPU
- cpu_tensor =gpu_tensor.cpu()
- # define pytorch tensors
- x =torch.randn(10,20)
- y =torch.ones(20,5)
- # `@` mean matrix multiplication from python3.5, PEP-0465
- res =x @y
- # get the shape
- res.shape # torch.Size([10, 5])
從張量到變量
張量是 PyTorch 的一個完美組件,但是要想構建神經網絡這還遠遠不夠。反向傳播怎么辦?當然,我們可以手動實現它,但是真的需要這樣做嗎?幸好還有自動微分。為了支持這個功能,PyTorch 提供了變量,它是張量之上的封裝。如此,我們可以構建自己的計算圖,并自動計算梯度。每個變量實例都有兩個屬性:包含初始張量本身的.data,以及包含相應張量梯度的.grad
- importtorch
- fromtorch.autograd importVariable
- # define an inputs
- x_tensor =torch.randn(10,20)
- y_tensor =torch.randn(10,5)
- x =Variable(x_tensor,requires_grad=False)
- y =Variable(y_tensor,requires_grad=False)
- # define some weights
- w =Variable(torch.randn(20,5),requires_grad=True)
- # get variable tensor
- print(type(w.data))# torch.FloatTensor
- # get variable gradient
- print(w.grad)# None
- loss =torch.mean((y -x @w)**2)
- # calculate the gradients
- loss.backward()
- print(w.grad)# some gradients
- # manually apply gradients
- w.data -=0.01*w.grad.data
- # manually zero gradients after update
- w.grad.data.zero_()
你也許注意到我們手動計算了自己的梯度,這樣看起來很麻煩,我們能使用優化器嗎?當然。
- importtorch
- fromtorch.autograd importVariable
- importtorch.nn.functional asF
- x =Variable(torch.randn(10,20),requires_grad=False)
- y =Variable(torch.randn(10,3),requires_grad=False)
- # define some weights
- w1 =Variable(torch.randn(20,5),requires_grad=True)
- w2 =Variable(torch.randn(5,3),requires_grad=True)
- learning_rate =0.1
- loss_fn =torch.nn.MSELoss()
- optimizer =torch.optim.SGD([w1,w2],lr=learning_rate)
- forstep inrange(5):
- pred =F.sigmoid(x @w1)
- pred =F.sigmoid(pred @w2)
- loss =loss_fn(pred,y)
- # manually zero all previous gradients
- optimizer.zero_grad()
- # calculate new gradients
- loss.backward()
- # apply new gradients
- optimizer.step()
并不是所有的變量都可以自動更新。但是你應該可以從最后一段代碼中看到重點:我們仍然需要在計算新梯度之前將它手動歸零。這是 PyTorch 的核心理念之一。有時我們會不太明白為什么要這么做,但另一方面,這樣可以讓我們充分控制自己的梯度。
靜態圖 vs 動態圖
PyTorch 和 TensorFlow 的另一個主要區別在于其不同的計算圖表現形式。TensorFlow 使用靜態圖,這意味著我們是先定義,然后不斷使用它。在 PyTorch 中,每次正向傳播都會定義一個新計算圖。在開始階段,兩者之間或許差別不是很大,但動態圖會在你希望調試代碼,或定義一些條件語句時顯現出自己的優勢。就像你可以使用自己最喜歡的 debugger 一樣!
你可以比較一下 while 循環語句的下兩種定義——第一個是 TensorFlow 中,第二個是 PyTorch 中:
- importtensorflow astf
- first_counter =tf.constant(0)
- second_counter =tf.constant(10)
- some_value =tf.Variable(15)
- # condition should handle all args:
- defcond(first_counter,second_counter,*args):
- returnfirst_counter <second_counter
- defbody(first_counter,second_counter,some_value):
- first_counter =tf.add(first_counter,2)
- second_counter =tf.add(second_counter,1)
- returnfirst_counter,second_counter,some_value
- c1,c2,val =tf.while_loop(
- cond,body,[first_counter,second_counter,some_value])
- withtf.Session()assess:
- sess.run(tf.global_variables_initializer())
- counter_1_res,counter_2_res =sess.run([c1,c2])
- importtorch
- first_counter =torch.Tensor([0])
- second_counter =torch.Tensor([10])
- some_value =torch.Tensor(15)
- while(first_counter <second_counter)[0]:
- first_counter +=2
- second_counter +=1
看起來第二種方法比第一個簡單多了,你覺得呢?
模型定義
現在我們看到,想在 PyTorch 中創建 if/else/while 復雜語句非常容易。不過讓我們先回到常見模型中,PyTorch 提供了非常類似于 Keras 的、即開即用的層構造函數:
神經網絡包(nn)定義了一系列的模塊,它可以粗略地等價于神經網絡的層。模塊接收輸入變量并計算輸出變量,但也可以保存內部狀態,例如包含可學習參數的變量。nn 包還定義了一組在訓練神經網絡時常用的損失函數。
- fromcollections importOrderedDict
- importtorch.nn asnn
- # Example of using Sequential
- model =nn.Sequential(
- nn.Conv2d(1,20,5),
- nn.ReLU(),
- nn.Conv2d(20,64,5),
- nn.ReLU()
- )
- # Example of using Sequential with OrderedDict
- model =nn.Sequential(OrderedDict([
- ('conv1',nn.Conv2d(1,20,5)),
- ('relu1',nn.ReLU()),
- ('conv2',nn.Conv2d(20,64,5)),
- ('relu2',nn.ReLU())
- ]))
- output =model(some_input)
如果你想要構建復雜的模型,我們可以將 nn.Module 類子類化。當然,這兩種方式也可以互相結合。
- fromtorch importnn
- classModel(nn.Module):
- def__init__(self):
- super().__init__()
- self.feature_extractor =nn.Sequential(
- nn.Conv2d(3,12,kernel_size=3,padding=1,stride=1),
- nn.Conv2d(12,24,kernel_size=3,padding=1,stride=1),
- )
- self.second_extractor =nn.Conv2d(
- 24,36,kernel_size=3,padding=1,stride=1)
- defforward(self,x):
- x =self.feature_extractor(x)
- x =self.second_extractor(x)
- # note that we may call same layer twice or mode
- x =self.second_extractor(x)
- returnx
在__init__方法中,我們需要定義之后需要使用的所有層。在正向方法中,我們需要提出如何使用已經定義的層的步驟。而在反向傳播上,和往常一樣,計算是自動進行的。
自定義層
如果我們想要定義一些非標準反向傳播模型要怎么辦?這里有一個例子——XNOR 網絡:
在這里我們不會深入細節,如果你對它感興趣,可以參考一下原始論文:
https://arxiv.org/abs/1603.05279
與我們問題相關的是反向傳播需要權重必須介于-1 到 1 之間。在 PyTorch 中,這可以很容易實現:
- importtorch
- classMyFunction(torch.autograd.Function):
- @staticmethod
- defforward(ctx,input):
- ctx.save_for_backward(input)
- output =torch.sign(input)
- returnoutput
- @staticmethod
- defbackward(ctx,grad_output):
- # saved tensors - tuple of tensors, so we need get first
- input,=ctx.saved_variables
- grad_output[input.ge(1)]=0
- grad_output[input.le(-1)]=0
- returngrad_output
- # usage
- x =torch.randn(10,20)
- y =MyFunction.apply(x)
- # or
- my_func =MyFunction.apply
- y =my_func(x)
- # and if we want to use inside nn.Module
- classMyFunctionModule(torch.nn.Module):
- defforward(self,x):
- returnMyFunction.apply(x)
正如你所見,我們應該只定義兩種方法:一個為正向傳播,一個為反向傳播。如果我們需要從正向通道訪問一些變量,我們可以將它們存儲在 ctx 變量中。注意:在此前的 API 正向/反向傳播不是靜態的,我們存儲變量需要以 self.save_for_backward(input) 的形式,并以 input, _ = self.saved_tensors 的方式接入。
在 CUDA 上訓練模型
我們曾經討論過傳遞一個張量到 CUDA 上。但如果希望傳遞整個模型,我們可以通過調用.cuda() 來完成,并將每個輸入變量傳遞到.cuda() 中。在所有計算后,我們需要用返回.cpu() 的方法來獲得結果。
同時,PyTorch 也支持在源代碼中直接分配設備:
- importtorch
- ### tensor example
- x_cpu =torch.randn(10,20)
- w_cpu =torch.randn(20,10)
- # direct transfer to the GPU
- x_gpu =x_cpu.cuda()
- w_gpu =w_cpu.cuda()
- result_gpu =x_gpu @w_gpu
- # get back from GPU to CPU
- result_cpu =result_gpu.cpu()
- ### model example
- modelmodel =model.cuda()
- # train step
- inputs =Variable(inputs.cuda())
- outputs =model(inputs)
- # get back from GPU to CPU
- outputsoutputs =outputs.cpu()
因為有些時候我們想在 CPU 和 GPU 中運行相同的模型,而無需改動代碼,我們會需要一種封裝:
- classTrainer:
- def__init__(self,model,use_cuda=False,gpu_idx=0):
- self.use_cuda =use_cuda
- self.gpu_idx =gpu_idx
- selfself.model =self.to_gpu(model)
- defto_gpu(self,tensor):
- ifself.use_cuda:
- returntensor.cuda(self.gpu_idx)
- else:
- returntensor
- deffrom_gpu(self,tensor):
- ifself.use_cuda:
- returntensor.cpu()
- else:
- returntensor
- deftrain(self,inputs):
- inputs =self.to_gpu(inputs)
- outputs =self.model(inputs)
- outputs =self.from_gpu(outputs)
權重初始化
在 TesnorFlow 中權重初始化主要是在張量聲明中進行的。PyTorch 則提供了另一種方法:首先聲明張量,隨后在下一步里改變張量的權重。權重可以用調用 torch.nn.init 包中的多種方法初始化為直接訪問張量的屬性。這個決定或許并不直接了當,但當你希望初始化具有某些相同初始化類型的層時,它就會變得有用。
- importtorch
- fromtorch.autograd importVariable
- # new way with `init` module
- w =torch.Tensor(3,5)
- torch.nn.init.normal(w)
- # work for Variables also
- w2 =Variable(w)
- torch.nn.init.normal(w2)
- # old styled direct access to tensors data attribute
- w2.data.normal_()
- # example for some module
- defweights_init(m):
- classname =m.__class__.__name__
- ifclassname.find('Conv')!=-1:
- m.weight.data.normal_(0.0,0.02)
- elifclassname.find('BatchNorm')!=-1:
- m.weight.data.normal_(1.0,0.02)
- m.bias.data.fill_(0)
- # for loop approach with direct access
- classMyModel(nn.Module):
- def__init__(self):
- form inself.modules():
- ifisinstance(m,nn.Conv2d):
- n =m.kernel_size[0]*m.kernel_size[1]*m.out_channels
- m.weight.data.normal_(0,math.sqrt(2./n))
- elifisinstance(m,nn.BatchNorm2d):
- m.weight.data.fill_(1)
- m.bias.data.zero_()
- elifisinstance(m,nn.Linear):
- m.bias.data.zero_()
反向排除子圖
有時,當你希望保留模型中的某些層或者為生產環境做準備的時候,禁用某些層的自動梯度機制非常有用。在這種思路下,PyTorch 設計了兩個 flag:requires_grad 和 volatile。第一個可以禁用當前層的梯度,但子節點仍然可以計算。第二個可以禁用自動梯度,同時效果沿用至所有子節點。
- importtorch
- fromtorch.autograd importVariable
- # requires grad
- # If there’s a single input to an operation that requires gradient,
- # its output will also require gradient.
- x =Variable(torch.randn(5,5))
- y =Variable(torch.randn(5,5))
- z =Variable(torch.randn(5,5),requires_grad=True)
- a =x +y
- a.requires_grad # False
- b =a +z
- b.requires_grad # True
- # Volatile differs from requires_grad in how the flag propagates.
- # If there’s even a single volatile input to an operation,
- # its output is also going to be volatile.
- x =Variable(torch.randn(5,5),requires_grad=True)
- y =Variable(torch.randn(5,5),volatile=True)
- a =x +y
- a.requires_grad # False
訓練過程
當然,PyTorch 還有一些其他賣點。例如你可以設定學習速率,讓它以特定規則進行變化。或者你可以通過簡單的訓練標記允許/禁止批規范層和 dropout。如果你想要做的話,讓 CPU 和 GPU 的隨機算子不同也是可以的。
- # scheduler example
- fromtorch.optim importlr_scheduler
- optimizer =torch.optim.SGD(model.parameters(),lr=0.01)
- scheduler =lr_scheduler.StepLR(optimizer,step_size=30,gamma=0.1)
- forepoch inrange(100):
- scheduler.step()
- train()
- validate()
- # Train flag can be updated with boolean
- # to disable dropout and batch norm learning
- model.train(True)
- # execute train step
- model.train(False)
- # run inference step
- # CPU seed
- torch.manual_seed(42)
- # GPU seed
- torch.cuda.manual_seed_all(42)
同時,你也可以添加模型信息,或存儲/加載一小段代碼。如果你的模型是由 OrderedDict 或基于類的模型字符串,它的表示會包含層名。
- fromcollections importOrderedDict
- importtorch.nn asnn
- model =nn.Sequential(OrderedDict([
- ('conv1',nn.Conv2d(1,20,5)),
- ('relu1',nn.ReLU()),
- ('conv2',nn.Conv2d(20,64,5)),
- ('relu2',nn.ReLU())
- ]))
- print(model)
- # Sequential (
- # (conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
- # (relu1): ReLU ()
- # (conv2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
- # (relu2): ReLU ()
- # )
- # save/load only the model parameters(prefered solution)
- torch.save(model.state_dict(),save_path)
- model.load_state_dict(torch.load(save_path))
- # save whole model
- torch.save(model,save_path)
- model =torch.load(save_path)
根據 PyTorch 文檔,用 state_dict() 的方式存儲文檔更好。
記錄
訓練過程的記錄是一個非常重要的部分。不幸的是,PyTorch 目前還沒有像 Tensorboard 這樣的東西。所以你只能使用普通文本記錄 Python 了,你也可以試試一些第三方庫:
- logger:https://github.com/oval-group/logger
- Crayon:https://github.com/torrvision/crayon
- tensorboard_logger:https://github.com/TeamHG-Memex/tensorboard_logger
- tensorboard-pytorch:https://github.com/lanpa/tensorboard-pytorch
- Visdom:https://github.com/facebookresearch/visdom
掌控數據
你可能會記得 TensorFlow 中的數據加載器,甚至想要實現它的一些功能。對于我來說,我花了四個小時來掌握其中所有管道的執行原理。
首先,我想在這里添加一些代碼,但我認為上圖足以解釋它的基礎理念了。
PyTorch 開發者不希望重新發明輪子,他們只是想要借鑒多重處理。為了構建自己的數據加載器,你可以從 torch.utils.data.Dataset 繼承類,并更改一些方法:
- importtorch
- importtorchvision astv
- classImagesDataset(torch.utils.data.Dataset):
- def__init__(self,df,transform=None,
- loader=tv.datasets.folder.default_loader):
- self.df =df
- self.transform =transform
- self.loader =loader
- def__getitem__(self,index):
- row =self.df.iloc[index]
- target =row['class_']
- path =row['path']
- img =self.loader(path)
- ifself.transform isnotNone:
- img =self.transform(img)
- returnimg,target
- def__len__(self):
- n,_ =self.df.shape
- returnn
- # what transformations should be done with our images
- data_transforms =tv.transforms.Compose([
- tv.transforms.RandomCrop((64,64),padding=4),
- tv.transforms.RandomHorizontalFlip(),
- tv.transforms.ToTensor(),
- ])
- train_df =pd.read_csv('path/to/some.csv')
- # initialize our dataset at first
- train_dataset =ImagesDataset(
- df=train_df,
- transform=data_transforms
- )
- # initialize data loader with required number of workers and other params
- train_loader =torch.utils.data.DataLoader(train_dataset,
- batch_size=10,
- shuffle=True,
- num_workers=16)
- # fetch the batch(call to `__getitem__` method)
- forimg,target intrain_loader:
- pass
有兩件事你需要事先知道:
1. PyTorch 的圖維度和 TensorFlow 的不同。前者的是 [Batch_size × channels × height × width] 的形式。但如果你沒有通過預處理步驟 torchvision.transforms.ToTensor() 進行交互,則可以進行轉換。在 transforms 包中還有很多有用小工具。
2. 你很可能會使用固定內存的 GPU,對此,你只需要對 cuda() 調用額外的標志 async = True,并從標記為 pin_memory = True 的 DataLoader 中獲取固定批次。
最終架構
現在我們了解了模型、優化器和很多其他細節。是時候來個總結了:
這里有一段用于解讀的偽代碼:
- classImagesDataset(torch.utils.data.Dataset):
- pass
- classNet(nn.Module):
- pass
- model =Net()
- optimizer =torch.optim.SGD(model.parameters(),lr=0.01)
- scheduler =lr_scheduler.StepLR(optimizer,step_size=30,gamma=0.1)
- criterion =torch.nn.MSELoss()
- dataset =ImagesDataset(path_to_images)
- data_loader =torch.utils.data.DataLoader(dataset,batch_size=10)
- train =True
- forepoch inrange(epochs):
- iftrain:
- lr_scheduler.step()
- forinputs,labels indata_loader:
- inputs =Variable(to_gpu(inputs))
- labels =Variable(to_gpu(labels))
- outputs =model(inputs)
- loss =criterion(outputs,labels)
- iftrain:
- optimizer.zero_grad()
- loss.backward()
- optimizer.step()
- ifnottrain:
- save_best_model(epoch_validation_accuracy)
結論
希望本文可以讓你了解 PyTorch 的如下特點:
- 它可以用來代替 Numpy
- 它的原型設計非常快
- 調試和使用條件流非常簡單
- 有很多方便且開箱即用的工具
PyTorch 是一個正在快速發展的框架,背靠一個富有活力的社區。現在是嘗試 PyTorch 的好時機。
【本文是51CTO專欄機構“機器之心”的原創譯文,微信公眾號“機器之心( id: almosthuman2014)”】