推荐学习b站视频
torch.Tensor()默认是torch.FloatTensor()的简称,创建的为float32位的数据类型;
torch.LongTensor是64位整型
optimizer.zero_grad() 清空过往梯度
loss.backward():反向传播
optimizer.step():在反向传播后更新参数
使用GPU加速,注意要把模型和数据都放到GPU上
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model=MyModel().to(device)
x,y=x.to(device),y.to(device)
import torch
import torch.nn as nn
import torch.utils.data as Data
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def fizzbuzz_encode(number):
if number%15==0:
return 3
elif number%5==0:
return 2
elif number%3==0:
return 1
return 0
def fizzbuzz_decode(number,label):
return [str(number),'fizz','buzz','fizzbuzz'][label]
NUM_DIGITS=10
def binary_encode(number):
return np.array([number >> d&1 for d in range(NUM_DIGITS)][::-1])
x_train=torch.Tensor([binary_encode(number) for number in range(100,1024)])
y_train=torch.LongTensor([fizzbuzz_encode(number) for number in range(100,1024)])
class MyDataset(Data.Dataset):
def __init__(self,x_train,y_train):
self.x_train=x_train
self.y_train=y_train
def __getitem__(self, idx):
return self.x_train[idx],self.y_train[idx]
def __len__(self):
return len(self.x_train)
train_dataset=MyDataset(x_train,y_train)
train_loader=Data.DataLoader(train_dataset,batch_size=16,shuffle=True)
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.liner1=nn.Linear(10,64)
self.activate1=nn.ReLU()
self.liner2=nn.Linear(64,8)
self.activate2=nn.ReLU()
self.liner3=nn.Linear(8,4)
def forward(self,x):
output=self.liner1(x)
output=self.activate1(output)
output=self.liner2(output)
output=self.activate2(output)
output=self.liner3(output)
return output
model=MyModel().to(device)
loss_fn=nn.CrossEntropyLoss()
optimizer=optim.Adam(model.parameters(),lr=0.01)
Epoch=100
for epoch in range(Epoch):
idx=0
for x,y in train_loader:
x,y=x.to(device),y.to(device)
pred=model(x)
loss=loss_fn(pred,y)
idx+=1
if idx%10==0:
print(loss)
optimizer.zero_grad()
loss.backward()
optimizer.step()
x_test=torch.Tensor([binary_encode(15)]).to(device)
pred=model(x_test)
softmax=nn.Softmax(dim=1)
pred=softmax(pred)
result=np.argmax(pred.cpu().detach().numpy(),1)
print(result)