cleanlab--cleanlab
Aditya Thyagarajan
109 行
3.9 KiB
Python
109 行
3.9 KiB
Python
# Copyright (C) 2017-2023 Cleanlab Inc.
|
|
# This file is part of cleanlab.
|
|
#
|
|
# cleanlab is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published
|
|
# by the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# cleanlab is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with cleanlab. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
"""
|
|
A PyTorch CNN which can be used for finding label issues in CIFAR-10 and CleanLearning with co-teaching.
|
|
|
|
Code adapted from: https://github.com/bhanML/Co-teaching/blob/master/model.py
|
|
|
|
You must have PyTorch installed: https://pytorch.org/get-started/locally/
|
|
"""
|
|
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def call_bn(bn, x):
|
|
return bn(x)
|
|
|
|
|
|
class CNN(nn.Module):
|
|
"""A CNN architecture shown to be a good baseline for a CIFAR-10 benchmark.
|
|
|
|
Parameters
|
|
----------
|
|
input_channel : int
|
|
n_outputs : int
|
|
dropout_rate : float
|
|
top_bn : bool
|
|
|
|
Methods
|
|
-------
|
|
forward
|
|
forward pass in PyTorch"""
|
|
|
|
def __init__(self, input_channel=3, n_outputs=10, dropout_rate=0.25, top_bn=False):
|
|
self.dropout_rate = dropout_rate
|
|
self.top_bn = top_bn
|
|
super(CNN, self).__init__()
|
|
self.c1 = nn.Conv2d(input_channel, 128, kernel_size=3, stride=1, padding=1)
|
|
self.c2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
|
|
self.c3 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
|
|
self.c4 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
|
|
self.c5 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
|
|
self.c6 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
|
|
self.c7 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=0)
|
|
self.c8 = nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=0)
|
|
self.c9 = nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=0)
|
|
self.l_c1 = nn.Linear(128, n_outputs)
|
|
self.bn1 = nn.BatchNorm2d(128)
|
|
self.bn2 = nn.BatchNorm2d(128)
|
|
self.bn3 = nn.BatchNorm2d(128)
|
|
self.bn4 = nn.BatchNorm2d(256)
|
|
self.bn5 = nn.BatchNorm2d(256)
|
|
self.bn6 = nn.BatchNorm2d(256)
|
|
self.bn7 = nn.BatchNorm2d(512)
|
|
self.bn8 = nn.BatchNorm2d(256)
|
|
self.bn9 = nn.BatchNorm2d(128)
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
):
|
|
h = x
|
|
h = self.c1(h)
|
|
h = F.leaky_relu(call_bn(self.bn1, h), negative_slope=0.01)
|
|
h = self.c2(h)
|
|
h = F.leaky_relu(call_bn(self.bn2, h), negative_slope=0.01)
|
|
h = self.c3(h)
|
|
h = F.leaky_relu(call_bn(self.bn3, h), negative_slope=0.01)
|
|
h = F.max_pool2d(h, kernel_size=2, stride=2)
|
|
h = F.dropout2d(h, p=self.dropout_rate)
|
|
|
|
h = self.c4(h)
|
|
h = F.leaky_relu(call_bn(self.bn4, h), negative_slope=0.01)
|
|
h = self.c5(h)
|
|
h = F.leaky_relu(call_bn(self.bn5, h), negative_slope=0.01)
|
|
h = self.c6(h)
|
|
h = F.leaky_relu(call_bn(self.bn6, h), negative_slope=0.01)
|
|
h = F.max_pool2d(h, kernel_size=2, stride=2)
|
|
h = F.dropout2d(h, p=self.dropout_rate)
|
|
|
|
h = self.c7(h)
|
|
h = F.leaky_relu(call_bn(self.bn7, h), negative_slope=0.01)
|
|
h = self.c8(h)
|
|
h = F.leaky_relu(call_bn(self.bn8, h), negative_slope=0.01)
|
|
h = self.c9(h)
|
|
h = F.leaky_relu(call_bn(self.bn9, h), negative_slope=0.01)
|
|
h = F.avg_pool2d(h, kernel_size=h.data.shape[2])
|
|
|
|
h = h.view(h.size(0), h.size(1))
|
|
logit = self.l_c1(h)
|
|
if self.top_bn:
|
|
logit = call_bn(self.bn_c1, logit)
|
|
return logit
|