使用深度学习进行疟疾检测 | PyTorch版

from torchvision.datasets import ImageFolderfrom torchvision.transforms import ToTensorfrom torchvision.transforms import Composefrom torchvision.transforms import Resizefrom torchvision.transforms import ToPILImage
dataset = ImageFolder(data_dir, transform=Compose([Resize((128, 128), interpolation=3),ToTensor()]))
classes = dataset.classesprint(classes)len(dataset)

import matplotlib.pyplot as plt
def show_example(img, label):    print('Label: ', dataset.classes[label], "("+str(label)+")")    plt.imshow(img.permute(1, 2, 0))

show_example(*dataset[2222])

show_example(*dataset[22269])

random_seed = 42torch.manual_seed(random_seed);

val_size = 6000train_size = 16000test_size = len(dataset) - val_size - train_size
train_ds, val_ds, test_ds = random_split(dataset, [train_size, val_size, test_size])len(train_ds), len(val_ds), len(test_ds)

from torch.utils.data.dataloader import DataLoaderbatch_size=128

train_dl = DataLoader(train_ds, batch_size, shuffle=True, num_workers=4, pin_memory=True)val_dl = DataLoader(val_ds, batch_size*2, num_workers=4, pin_memory=True)

from torchvision.utils import make_grid
def show_batch(dl):    for images, labels in dl:        fig, ax = plt.subplots(figsize=(12, 6))        ax.set_xticks([]); ax.set_yticks([])        ax.imshow(make_grid(images, nrow=16).permute(1, 2, 0))        break

def accuracy(outputs, labels):    _, preds = torch.max(outputs, dim=1)    return torch.tensor(torch.sum(preds == labels).item() / len(preds))
class ImageClassificationBase(nn.Module):    def training_step(self, batch):        images, labels = batch         out = self(images)                  # Generate predictions        loss = F.cross_entropy(out, labels) # Calculate loss        return loss        def validation_step(self, batch):        images, labels = batch         out = self(images)                    # Generate predictions        loss = F.cross_entropy(out, labels)   # Calculate loss        acc = accuracy(out, labels)           # Calculate accuracy        return {'val_loss': loss.detach(), 'val_acc': acc}            def validation_epoch_end(self, outputs):        batch_losses = [x['val_loss'] for x in outputs]        epoch_loss = torch.stack(batch_losses).mean()   # Combine losses        batch_accs = [x['val_acc'] for x in outputs]        epoch_acc = torch.stack(batch_accs).mean()      # Combine accuracies        return {'val_loss': epoch_loss.item(), 'val_acc': epoch_acc.item()}        def epoch_end(self, epoch, result):        print("Epoch [{}], train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(            epoch, result['train_loss'], result['val_loss'], result['val_acc']))

class Malaria2CnnModel(ImageClassificationBase):    def __init__(self):        super().__init__()        self.network = nn.Sequential(            nn.Conv2d(3, 32, kernel_size=3, padding=1),            nn.ReLU(),            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),            nn.ReLU(),            nn.MaxPool2d(2, 2), # output: 64 x 64 x 64
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),            nn.ReLU(),            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),            nn.ReLU(),            nn.MaxPool2d(2, 2), # output: 128 x 32 x 32
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),            nn.ReLU(),            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),            nn.ReLU(),            nn.MaxPool2d(2, 2), # output: 256 x 16 x 16
            nn.Flatten(),             nn.Linear(256*16*16, 1024),            nn.ReLU(),            nn.Linear(1024, 512),            nn.ReLU(),            nn.Linear(512, 2))            def forward(self, xb):        return self.network(xb)

model = Malaria2CnnModel()model

def get_default_device():    """Pick GPU if available, else CPU"""    if torch.cuda.is_available():        return torch.device('cuda')    else:        return torch.device('cpu')    def to_device(data, device):    """Move tensor(s) to chosen device"""    if isinstance(data, (list,tuple)):        return [to_device(x, device) for x in data]    return data.to(device, non_blocking=True)
class DeviceDataLoader():    """Wrap a dataloader to move data to a device"""    def __init__(self, dl, device):        self.dl = dl        self.device = device            def __iter__(self):        """Yield a batch of data after moving it to device"""        for b in self.dl:             yield to_device(b, self.device)
    def __len__(self):        """Number of batches"""        return len(self.dl)

device = get_default_device()device

train_dl = DeviceDataLoader(train_dl, device)val_dl = DeviceDataLoader(val_dl, device)to_device(model, device);

@torch.no_grad()def evaluate(model, val_loader):    model.eval()    outputs = [model.validation_step(batch) for batch in val_loader]    return model.validation_epoch_end(outputs)
def fit(epochs, lr, model, train_loader, val_loader, opt_func=torch.optim.SGD):    history = []    optimizer = opt_func(model.parameters(), lr)    for epoch in range(epochs):        # Training Phase         model.train()        train_losses = []        for batch in train_loader:            loss = model.training_step(batch)            train_losses.append(loss)            loss.backward()            optimizer.step()            optimizer.zero_grad()        # Validation phase        result = evaluate(model, val_loader)        result['train_loss'] = torch.stack(train_losses).mean().item()        model.epoch_end(epoch, result)        history.append(result)    return history

model = to_device(Malaria2CnnModel(), device)

evaluate(model, val_dl)

num_epochs = 10opt_func = torch.optim.Adamlr = 0.001

def plot_accuracies(history):    accuracies = [x['val_acc'] for x in history]    plt.plot(accuracies, '-x')    plt.xlabel('epoch')    plt.ylabel('accuracy')    plt.title('Accuracy vs. No. of epochs');

def plot_losses(history):    train_losses = [x.get('train_loss') for x in history]    val_losses = [x['val_loss'] for x in history]    plt.plot(train_losses, '-bx')    plt.plot(val_losses, '-rx')    plt.xlabel('epoch')    plt.ylabel('loss')    plt.legend(['Training', 'Validation'])    plt.title('Loss vs. No. of epochs');

history = fit(num_epochs, lr, model, train_dl, val_dl, opt_func)

plot_accuracies(history)

plot_losses(history)

def predict_image(img, model):    # Convert to a batch of 1    xb = to_device(img.unsqueeze(0), device)    # Get predictions from model    yb = model(xb)    # Pick index with highest probability    _, preds  = torch.max(yb, dim=1)    # Retrieve the class label    return dataset.classes[preds[0].item()]

img, label = test_ds[0]plt.imshow(img.permute(1, 2, 0))print('Label:', dataset.classes[label], ', Predicted:', predict_image(img, model))

test_loader = DeviceDataLoader(DataLoader(test_ds, batch_size*2), device)result = evaluate(model, test_loader)result