def main():
train_dataset = MNIST(root='./data',
train=True,
download=True,
transform=transforms.ToTensor())
test_dataset = MNIST(root='./data',
train=False,
download=True,
transform=transforms.ToTensor())
train_loader = DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=2)
test_loader = DataLoader(test_dataset,
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=2)
net = Network(1, 64, 5, 10)
if USE_CUDA:
net = net.cuda()
opt = optim.SGD(net.parameters(),
lr=LEARNING_RATE,
weight_decay=WEIGHT_DECAY,
momentum=.9,
nesterov=True)
if not os.path.exists('checkpoint'):
os.mkdir('checkpoint')
for epoch in range(1, EPOCHS + 1):
print('[Epoch %d]' % epoch)
train_loss = 0
train_correct, train_total = 0, 0
start_point = time.time()
for inputs, labels in train_loader:
inputs, labels = Variable(inputs), Variable(labels)
if USE_CUDA:
inputs, labels = inputs.cuda(), labels.cuda()
opt.zero_grad()
preds = F.log_softmax(net(inputs), dim=1)
loss = F.cross_entropy(preds, labels)
loss.backward()
opt.step()
train_loss += loss.item()
train_correct += (preds.argmax(dim=1) == labels).sum().item()
train_total += len(preds)
print('train-acc : %.4f%% train-loss : %.5f' %
(100 * train_correct / train_total,
train_loss / len(train_loader)))
print('elapsed time: %ds' % (time.time() - start_point))
test_loss = 0
test_correct, test_total = 0, 0
for inputs, labels in test_loader:
with torch.no_grad():
inputs, labels = Variable(inputs), Variable(labels)
if USE_CUDA:
inputs, labels = inputs.cuda(), labels.cuda()
preds = F.softmax(net(inputs), dim=1)
test_loss += F.cross_entropy(preds, labels).item()
test_correct += (preds.argmax(dim=1) == labels).sum().item()
test_total += len(preds)
print('test-acc : %.4f%% test-loss : %.5f' %
(100 * test_correct / test_total, test_loss / len(test_loader)))
torch.save(net.state_dict(),
'./checkpoint/checkpoint-%04d.bin' % epoch)
cfg.TRAIN.score_com_strength,
cfg.TRAIN.scale_com_strength,
cfg.TRAIN.NMS_THRESH,
cfg.TRAIN.NMS_KSIZE,
cfg.TRAIN.TOPK,
cfg.MODEL.GAUSSIAN_KSIZE,
cfg.MODEL.GAUSSIAN_SIGMA,
cfg.MODEL.KSIZE,
cfg.MODEL.padding,
cfg.MODEL.dilation,
cfg.MODEL.scale_list,
)
des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH)
model = Network(det, des, cfg.LOSS.SCORE, cfg.LOSS.PAIR, cfg.PATCH.SIZE,
cfg.TRAIN.TOPK)
model.cuda()
from thop import profile, clever_format
flop, para = profile(det, (torch.randn((1, 1, 640, 480), device=device), ))
'''
flop
364512000.0
para
1377.0
'''
flop2, para2 = profile(des, (torch.randn((1000, 1, 32, 32), device=device), ))
'''
flop2
16663040000.0
para2
