參考莫凡部落格進行mnist資料集的訓練，臨時記錄所使用的**。
import torch
import torch.nn as nn
import torch.utils.data as data
import torchvision
import matplotlib.pyplot as plt
torch.manual_seed(1)
epoch =
1batch_size =
50lr =
0.001
download_mnist =
true
train_data = torchvision.datasets.mnist(
root=
'./mnist'
, train=
true
, transform=torchvision.transforms.totensor(),
download=download_mnist,
)test_data = torchvision.datasets.mnist(root=
'./mnist'
, train=
false
)# 批處理
train_loader = data.dataloader(dataset=train_data, batch_size=batch_size, shuffle=
true
)# 測試
test_x = torch.unsqueeze(test_data.test_data, dim=1)
.type
(torch.floattensor)[:
2000]/
255test_y = test_data.test_labels[
:2000
]test_x = test_x.cuda(
)test_y = test_y.cuda(
)# 卷積(conv2d) -> 激勵函式(relu) -> 池化, 向下取樣 (maxpooling) ->
# 再來一遍 -> 展平多維的卷積成的特徵圖 -> 接入全連線層 (linear) -> 輸出
class
cnn(nn.module)
:def
__init__
(self)
:super
(cnn, self)
.__init__(
) self.conv1 = nn.sequential(
# 1x28x28
nn.conv2d(
in_channels=1,
out_channels=16,
kernel_size=5,
stride=1,
padding=2,
),# 16x28x28
nn.relu(),
nn.maxpool2d(kernel_size=2)
) self.conv2 = nn.sequential(
# 16x14x14
nn.conv2d(16,
32,5,
1,2)
,# 32x14x14
nn.relu(),
nn.maxpool2d(2)
,# 32x7x7
) self.out = nn.linear(32*
7*7,
10)defforward
(self, x)
: x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size(0)
,-1)
# 展平多維的卷積圖成 (batch_size, 32 * 7 * 7)\
output = self.out(x)
return output
cnn = cnn(
)cnn = cnn.cuda(
)print
(cnn)
optimizer = torch.optim.adam(cnn.parameters(
), lr=lr)
loss_func = nn.crossentropyloss(
)for epoch in
range
(epoch)
:for step,
(b_x, b_y)
inenumerate
(train_loader)
: output = cnn(b_x.cuda())
loss = loss_func(output, b_y.cuda())
optimizer.zero_grad(
) loss.backward(
) optimizer.step(
)test_output = cnn(test_x[:10
])pred_y = torch.
max(test_output,1)
[1].data.cpu(
).numpy(
).squeeze(
)print
(pred_y,
'prediction number'
)print
(test_y[:10
].cpu(
).numpy(),
'real number'
)
Pytorch中多GPU訓練
參考 在資料越來越多的時代，隨著模型規模引數的增多，以及資料量的不斷提公升，使用多gpu去訓練是不可避免的事情。pytorch在0.4.0及以後的版本中已經提供了多gpu訓練的方式，本文簡單講解下使用pytorch多gpu訓練的方式以及一些注意的地方。這裡我們談論的是單主機多gpus訓練，與分布式訓...
pytorch 多GPU訓練注意事項
1.多gpu訓練記得dataloader dataset dataset train,batch size config train batch shuffle config train shuffle num workers config train workers drop last true ...
使用指定GPU跑Pytorch
有如下兩種方法來指定需要使用的gpu。類似tensorflow指定gpu的方式，使用cuda visible devices。1.1 直接終端中設定 cuda visible devices 1 python my script.py 1.2 python 中設定 import os os.envi...