def__init__(self, root_dir, transform=None): """ Args: root_dir (string): Directory with all the images. transform (callable, optional): Optional transform to be applied on a sample. """ self.root_dir = root_dir self.transform = transform self.data_list = [] for i inrange(10): img_paths = glob.glob(root_dir + "%d/*.jpg" %i) for img_path in img_paths: img_label = {'img_path':img_path, 'label':i } self.data_list.append(img_label)
defforward(self, x): out = self.layer1(x) out = self.layer2(out) out = out.reshape(out.size(0), -1) out = self.fc(out) return F.log_softmax(out, dim=1)
model = ConvNet(num_classes=10).to(device) optimizer = optim.SGD(model.parameters(), lr=0.01)
在编写完搭建卷积神经网络的模型代码后,我们就可以开始对模型进行训练和对参数进行优化了。
# Define training loops for epoch inrange(20): model.train() loss_value = 0. acc_value = 0. num_batch = 0 with tqdm(total=len(train_loader), desc="Epoch %2d/20" %(epoch+1)) as pbar: for data, target in train_loader: data, target = data.to(device), target.to(device)
loss_value += loss.item() # get the index of the max log-probability pred = output.max(1, keepdim=True)[1] correct = pred.eq(target.view_as(pred)).sum().item() acc_value += correct / len(target)