train model with gpu

.gitignore

@@ -0,0 +1,3 @@
+data/
+__pycache__/
+*.pth

classification/train.py

@@ -1,10 +1,13 @@
+import os
 from resnet import model
 from cifar_data import CIFAR10
 
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.optim import lr_scheduler
+import torch.backends.cudnn as cudnn
+
 import torchvision
 import torchvision.transforms as transforms
 
@@ -31,8 +34,8 @@
 	transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
 ])
 
-dataset_train = CIFAR10(root='./data', train=True, download=True, transform=transforms_train)
-dataset_test = CIFAR10(root='./data', train=False, download=True, transform=transforms_test)
+dataset_train = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transforms_train)
+dataset_test = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms_test)
 train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=args.batch_size, 
 	                      shuffle=True, num_workers=args.num_worker)
 test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=100, 
@@ -46,7 +49,7 @@
 net = model.resnet18()
 net = net.to(device)
 if device == 'cuda':
-	net = cuda.nn.DataParallel(net)
+	net = torch.nn.DataParallel(net)
 	cudnn.benchmark = True
 
 
@@ -126,4 +129,4 @@ def test(epoch, best_acc):
 	best_acc = 0
 	for epoch in range(200):
 		train(epoch)
-		best_acc = test(epoch, best_acc)
+		best_acc = test(epoch, best_acc)