神經(jīng)網(wǎng)絡(luò)一般用GPU來(lái)跑，我們的神經(jīng)網(wǎng)絡(luò)框架一般也都安裝的GPU版本，本文就簡(jiǎn)單記錄一下GPU使用的編寫。

GPU的設(shè)置不在model，而是在Train的初始化上。

第一步是查看是否可以使用GPU

self.GPU_IN_USE = torch.cuda.is_available()

就是返回這個(gè)可不可以用GPU的函數(shù)，當(dāng)你的pytorch是cpu版本的時(shí)候，他就會(huì)返回False。

然后是：

self.device = torch.device('cuda' if self.GPU_IN_USE else 'cpu')

torch.device是代表將torch.tensor分配到哪個(gè)設(shè)備的函數(shù)

接著是，我看到了一篇文章，原來(lái)就是將網(wǎng)絡(luò)啊、數(shù)據(jù)啊、隨機(jī)種子啊、損失函數(shù)啊、等等等等直接轉(zhuǎn)移到CUDA上就好了！

于是下面就好理解多了：

轉(zhuǎn)移模型：

self.model = Net(num_channels=1, upscale_factor=self.upscale_factor, base_channel=64, num_residuals=4).to(self.device)

設(shè)置cuda的隨機(jī)種子：

torch.cuda.manual_seed(self.seed)

轉(zhuǎn)移損失函數(shù)：

self.criterion.cuda()

轉(zhuǎn)移數(shù)據(jù)：

data, target = data.to(self.device), target.to(self.device)

pytorch 網(wǎng)絡(luò)定義參數(shù)的后面無(wú)法加.cuda()

pytorch定義網(wǎng)絡(luò)__init__()的時(shí)候，參數(shù)不能加“cuda()", 不然參數(shù)不包含在state_dict()中，比如下面這種寫法是錯(cuò)誤的

self.W1 = nn.Parameter(torch.FloatTensor(3,3), requires_grad=True).cuda()

應(yīng)該去掉".cuda()"

self.W1 = nn.Parameter(torch.FloatTensor(3,3), requires_grad=True)

以上為個(gè)人經(jīng)驗(yàn)，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。