2022-01-10 16:05:23 +08:00
|
|
|
|
import os
|
|
|
|
|
import torch
|
|
|
|
|
import torch.nn as nn
|
|
|
|
|
import torch.nn.functional as F
|
|
|
|
|
import torch.optim as optim
|
|
|
|
|
import torchvision
|
|
|
|
|
from torchvision import datasets, transforms
|
|
|
|
|
import torchvision.models as models
|
|
|
|
|
from torch.cuda.amp import autocast as autocast
|
|
|
|
|
from torch.utils.data import Dataset, DataLoader
|
|
|
|
|
import numpy as np
|
|
|
|
|
import random
|
|
|
|
|
import cv2
|
|
|
|
|
from torch.profiler import profile, record_function, ProfilerActivity
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
batchsize = 12
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ManualDataset(Dataset):
|
|
|
|
|
def __init__(self, transform=None):
|
|
|
|
|
self.transform = transform
|
|
|
|
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
|
|
|
|
|
|
|
# data = np.random.random_integers(0,255,(28,28))
|
|
|
|
|
data = np.zeros((28, 28), dtype="float32")
|
|
|
|
|
data = data.astype("float32")
|
|
|
|
|
|
|
|
|
|
radiu = int(random.random()*7+3)
|
|
|
|
|
basex = random.randint(radiu, 29-radiu)
|
|
|
|
|
basey = random.randint(radiu, 29-radiu)
|
|
|
|
|
ifcircle = random.random()
|
|
|
|
|
|
|
|
|
|
if ifcircle > 0.5:
|
|
|
|
|
cv2.circle(data, (basex, basey), radiu, 255, -1)
|
|
|
|
|
angle = random.random() * 360
|
|
|
|
|
M = cv2.getRotationMatrix2D((basex, basey), angle, 1.0)
|
|
|
|
|
data = cv2.warpAffine(data, M, (28, 28))
|
|
|
|
|
label = 0
|
|
|
|
|
else:
|
|
|
|
|
cv2.rectangle(data, (basex-radiu, basey-radiu),
|
|
|
|
|
(basex+radiu, basey+radiu), 255, -1)
|
|
|
|
|
angle = random.random() * 360
|
|
|
|
|
M = cv2.getRotationMatrix2D((basex, basey), angle, 1.0)
|
|
|
|
|
data = cv2.warpAffine(data, M, (28, 28))
|
|
|
|
|
label = 1
|
|
|
|
|
|
|
|
|
|
# cv2.imwrite("test.jpg",data)
|
|
|
|
|
|
|
|
|
|
data = (data - 128) / 256.0
|
|
|
|
|
img = torch.from_numpy(data)
|
|
|
|
|
img = img.view(1, 28, 28)
|
|
|
|
|
return img, label
|
|
|
|
|
|
|
|
|
|
def __len__(self):
|
|
|
|
|
return 10000
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
train_loader = torch.utils.data.DataLoader(
|
|
|
|
|
ManualDataset(transform=transforms.Compose([
|
|
|
|
|
# transforms.ColorJitter(0.2,0.2),
|
|
|
|
|
# transforms.RandomRotation(30),
|
|
|
|
|
# transforms.RandomResizedCrop(28),
|
|
|
|
|
transforms.ToTensor(),
|
|
|
|
|
transforms.Normalize((128,), (256,)), ])),
|
|
|
|
|
batch_size=batchsize, shuffle=True, num_workers=2)
|
|
|
|
|
test_loader = torch.utils.data.DataLoader(
|
|
|
|
|
ManualDataset(transform=transforms.Compose([
|
|
|
|
|
transforms.ToTensor(),
|
|
|
|
|
transforms.Normalize((128,), (256,))])),
|
|
|
|
|
batch_size=batchsize, shuffle=True, num_workers=2)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Net(nn.Module):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
super(Net, self).__init__()
|
|
|
|
|
layers = []
|
|
|
|
|
layers += [nn.Conv2d(1, 8, kernel_size=5),
|
|
|
|
|
nn.MaxPool2d(kernel_size=2, stride=2), nn.Sigmoid()]
|
|
|
|
|
layers += [nn.Conv2d(8, 8, kernel_size=3),
|
|
|
|
|
nn.MaxPool2d(kernel_size=2, stride=2), nn.Sigmoid()]
|
|
|
|
|
layers += [nn.Conv2d(8, 10, kernel_size=5)]
|
|
|
|
|
self.features = nn.Sequential(*layers)
|
|
|
|
|
|
|
|
|
|
# self.conv1 = nn.Conv2d(1, 8, kernel_size=5)
|
|
|
|
|
# self.conv2 = self.__conv(8, 8, kernel_size=3)
|
|
|
|
|
# self.conv3 = self.__conv(8, 10, kernel_size=5)
|
|
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
# x = F.sigmoid(F.max_pool2d(self.conv1(x), 2))
|
|
|
|
|
# x = F.sigmoid(F.max_pool2d(self.conv2(x), 2))
|
|
|
|
|
# x = self.conv3(x)
|
|
|
|
|
x = self.features(x)
|
|
|
|
|
x = x.view(-1, 1*10)
|
|
|
|
|
|
|
|
|
|
return F.log_softmax(x, dim=1)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
model = Net()
|
2022-01-14 17:16:26 +08:00
|
|
|
|
model = model.cuda()
|
2022-01-10 16:05:23 +08:00
|
|
|
|
|
|
|
|
|
optimizer = optim.SGD(model.parameters(), lr=0.0005)
|
|
|
|
|
|
|
|
|
|
for batch_idx, (data, target) in enumerate(train_loader):
|
2022-01-14 17:16:26 +08:00
|
|
|
|
data = data.cuda()
|
|
|
|
|
target = target.cuda()
|
2022-01-10 16:05:23 +08:00
|
|
|
|
optimizer.zero_grad()
|
|
|
|
|
|
2022-01-14 17:16:26 +08:00
|
|
|
|
# with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_stack=True, record_shapes=True) as prof:
|
|
|
|
|
# with record_function("model_inference"):
|
|
|
|
|
# with autocast():
|
|
|
|
|
# output = model(data)
|
|
|
|
|
# loss = F.nll_loss(output, target)
|
|
|
|
|
# # print(prof.key_averages().table(sort_by="cpu_time_total"))
|
2022-01-10 16:05:23 +08:00
|
|
|
|
# print(prof.key_averages().table(sort_by="cpu_time_total"))
|
2022-01-14 17:16:26 +08:00
|
|
|
|
# prof.export_chrome_trace("trace.json")
|
2022-01-10 16:05:23 +08:00
|
|
|
|
|
2022-01-14 17:16:26 +08:00
|
|
|
|
with autocast():
|
|
|
|
|
output = model(data)
|
|
|
|
|
loss = F.nll_loss(output, target)
|
|
|
|
|
|
|
|
|
|
# # 反向传播在autocast上下文之外
|
|
|
|
|
# loss.backward()
|
|
|
|
|
# optimizer.step()
|
2022-01-10 16:05:23 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# from torch.cuda.amp import autocast as autocast
|
|
|
|
|
|
|
|
|
|
# # 创建model,默认是torch.FloatTensor
|
|
|
|
|
# model = Net().cuda()
|
|
|
|
|
# optimizer = optim.SGD(model.parameters(), ...)
|
|
|
|
|
|
|
|
|
|
# # 在训练最开始之前实例化一个GradScaler对象
|
|
|
|
|
# scaler = GradScaler()
|
|
|
|
|
|
|
|
|
|
# for epoch in epochs:
|
|
|
|
|
# for input, target in data:
|
|
|
|
|
# optimizer.zero_grad()
|
|
|
|
|
|
|
|
|
|
# # 前向过程(model + loss)开启 autocast
|
|
|
|
|
# with autocast():
|
|
|
|
|
# output = model(input)
|
|
|
|
|
# loss = loss_fn(output, target)
|
|
|
|
|
|
|
|
|
|
# # Scales loss. 为了梯度放大.
|
|
|
|
|
# scaler.scale(loss).backward()
|
|
|
|
|
|
|
|
|
|
# # scaler.step() 首先把梯度的值unscale回来.
|
|
|
|
|
# # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
|
|
|
|
|
# # 否则,忽略step调用,从而保证权重不更新(不被破坏)
|
|
|
|
|
# scaler.step(optimizer)
|
|
|
|
|
|
|
|
|
|
# # 准备着,查看是否要增大scaler
|
|
|
|
|
# scaler.update()
|