手写CNN.zip

import os import cv2 import sys import re import time import numpy as np from tqdm import tqdm from datetime import datetime from tabulate import tabulate import torch import torch.nn as nn from torch.utils.data.sampler import SubsetRandomSampler from torch.utils.data import Dataset import pandas as pd from torch.utils.tensorboard import SummaryWriter from A6_submission import Classifier, Params class NotMNIST_RGB(Dataset): def __init__(self, fname='train_data_rgb.npz'): train_data = np.load(fname, allow_pickle=True) """_train_images is a 4D array of size n_images x 28 x 28 x 3""" self._train_images, self._train_labels = train_data['images'], train_data[ 'labels'] # type: np.ndarray, np.ndarray """ pixel values converted to floating-point numbers and normalized to be between 0 and 1 to make them suitable for processing in CNNs """ self._train_images = self._train_images.astype(np.float32) / 255.0 """switch images from n_images x 28 x 28 x 3 to n_images x 3 x 28 x 28 since CNNs expect the channels to be the first dimension""" self._train_images = np.transpose(self._train_images, (0, 3, 1, 2)) self._train_labels = self._train_labels.astype(np.int64) self._n_train = self._train_images.shape[0] def __len__(self): return self._n_train def __getitem__(self, idx): assert idx < self._n_train, "Invalid idx: {} for n_train: {}".format(idx, self._n_train) images = self._train_images[idx, ...] labels = self._train_labels[idx] return images, labels def train(classifier, dataloader, criterion, optimizer, device): total_loss = 0 train_total = 0 train_correct = 0 n_batches = 0 # set CNN to training mode classifier.train() for batch_idx, (inputs, targets) in tqdm(enumerate(dataloader)): inputs = inputs.to(device) targets = targets.to(device) optimizer.zero_grad() outputs = classifier(inputs) loss = criterion(outputs, targets) total_loss += loss.item() loss.backward() optimizer.step() _, predicted = outputs.max(1) train_total += targets.size(0) train_correct += predicted.eq(targets).sum().item() n_batches += 1 mean_loss = total_loss / n_batches train_acc = 100. * train_correct / train_total return mean_loss, train_acc def evaluate(classifier, dataloader, criterion, vis, writer, iteration, device): total_loss = 0 _psnr_sum = 0 total_images = 0 correct_images = 0 n_batches = 0 _pause = 1 # set CNN to evaluation mode classifier.eval() total_test_time = 0 # disable gradients computation with torch.no_grad(): for batch_idx, (inputs, targets) in enumerate(dataloader): inputs = inputs.to(device) targets = targets.to(device) start_t = time.time() outputs = classifier(inputs) end_t = time.time() test_time = end_t - start_t total_test_time += test_time loss = criterion(outputs, targets) total_loss += loss.item() _, predicted = outputs.max(1) total_images += targets.size(0) is_correct = predicted.eq(targets) correct_images += is_correct.sum().item() n_batches += 1 if vis: inputs_np = inputs.detach().cpu().numpy() concat_imgs = [] for i in range(dataloader.batch_size): input_img = inputs_np[i, ...].squeeze() """switch image from 3 x 28 x 28 to 28 x 28 x 3 since opencv expects channels to be on the last axis copy is needed to resolve an opencv issue with memory layout: https://stackoverflow.com/questions/23830618/python-opencv-typeerror-layout-of-the-output-array -incompatible-with-cvmat """ input_img = np.transpose(input_img, (1, 2, 0)).copy() target = targets[i] output = predicted[i] _is_correct = is_correct[i].item() if _is_correct: col = (0, 1, 0) else: col = (0, 0, 1) pred_img = np.zeros_like(input_img) _text = '{}'.format(chr(65 + int(output))) cv2.putText(pred_img, _text, (8, 20), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, col, 1, cv2.LINE_AA) label_img = np.zeros_like(input_img) _text = '{}'.format(chr(65 + int(target))) cv2.putText(label_img, _text, (8, 20), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (1, 1, 1), 1, cv2.LINE_AA) concat_img = np.concatenate((input_img, label_img, pred_img), axis=0) concat_imgs.append(concat_img) vis_img = np.concatenate(concat_imgs, axis=1) if writer is not None: vis_img_uint8 = (vis_img * 255.0).astype(np.uint8) vis_img_tb = cv2.cvtColor(vis_img_uint8, cv2.COLOR_BGR2RGB) """tensorboard expects channels in the first axis""" vis_img_tb = np.transpose(vis_img_tb, axes=[2, 0, 1]) writer.add_image('evaluation', vis_img_tb, iteration) if vis == 2: cv2.imshow('Press Esc to exit, Space to resume, any other key for next batch', vis_img) k = cv2.waitKey(1 - _pause) if k == 27: sys.exit(0) elif k == ord('q'): vis = 0 cv2.destroyWindow('concat_imgs') break elif k == 32: _pause = 1 - _pause mean_loss = total_loss / n_batches acc = 100. * float(correct_images) / float(total_images) test_speed = float(total_images) / total_test_time if vis == 2: cv2.destroyWindow('concat_imgs') return mean_loss, acc, total_images, test_speed def main(): params = Params() # optional command line argument parsing try: import paramparse except ImportError: pass else: paramparse.process(params) # init device if params.use_gpu and torch.cuda.is_available(): device = torch.device("cuda") print('Running on GPU: {}'.format(torch.cuda.get_device_name(0))) else: device = torch.device("cpu") print('Running on CPU') # load dataset train_set = NotMNIST_RGB() num_train = len(train_set) indices = list(range(num_train)) assert params.val_ratio > 0, "Zero validation ratio is not allowed " split = int(np.floor((1.0 - params.val_ratio) * num_train)) train_idx, val_idx = indices[:split], indices[split:] n_train = len(train_idx) n_val = len(val_idx) print('Training samples: {}\n' 'Validation samples: {}\n' ''.format( n_train, n_val, )) train_sampler = SubsetRandomSampler(train_idx) val_sampler = SubsetRandomSampler(val_idx) train_dataloader = torch.utils.data.DataLoader(train_set, batch_size=params.train_batch_size, sampler=train_sampler, num_workers=params.n_workers) val_dataloader = torch.utils.data.DataLoader(train_set, batch_size=params.val_batch_size,