-
Notifications
You must be signed in to change notification settings - Fork 20
/
jocor.py
160 lines (125 loc) · 5.82 KB
/
jocor.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
# -*- coding:utf-8 -*-
import torch
import torch.nn.functional as F
from torch.autograd import Variable
from model.cnn import MLPNet,CNN
import numpy as np
from common.utils import accuracy
from algorithm.loss import loss_jocor
class JoCoR:
def __init__(self, args, train_dataset, device, input_channel, num_classes):
# Hyper Parameters
self.batch_size = 128
learning_rate = args.lr
if args.forget_rate is None:
if args.noise_type == "asymmetric":
forget_rate = args.noise_rate / 2
else:
forget_rate = args.noise_rate
else:
forget_rate = args.forget_rate
self.noise_or_not = train_dataset.noise_or_not
# Adjust learning rate and betas for Adam Optimizer
mom1 = 0.9
mom2 = 0.1
self.alpha_plan = [learning_rate] * args.n_epoch
self.beta1_plan = [mom1] * args.n_epoch
for i in range(args.epoch_decay_start, args.n_epoch):
self.alpha_plan[i] = float(args.n_epoch - i) / (args.n_epoch - args.epoch_decay_start) * learning_rate
self.beta1_plan[i] = mom2
# define drop rate schedule
self.rate_schedule = np.ones(args.n_epoch) * forget_rate
self.rate_schedule[:args.num_gradual] = np.linspace(0, forget_rate ** args.exponent, args.num_gradual)
self.device = device
self.num_iter_per_epoch = args.num_iter_per_epoch
self.print_freq = args.print_freq
self.co_lambda = args.co_lambda
self.n_epoch = args.n_epoch
self.train_dataset = train_dataset
if args.model_type == "cnn":
self.model1 = CNN(input_channel=input_channel, n_outputs=num_classes)
self.model2 = CNN(input_channel=input_channel, n_outputs=num_classes)
elif args.model_type == "mlp":
self.model1 = MLPNet()
self.model2 = MLPNet()
self.model1.to(device)
print(self.model1.parameters)
self.model2.to(device)
print(self.model2.parameters)
self.optimizer = torch.optim.Adam(list(self.model1.parameters()) + list(self.model2.parameters()),
lr=learning_rate)
self.loss_fn = loss_jocor
self.adjust_lr = args.adjust_lr
# Evaluate the Model
def evaluate(self, test_loader):
print('Evaluating ...')
self.model1.eval() # Change model to 'eval' mode.
self.model2.eval() # Change model to 'eval' mode
correct1 = 0
total1 = 0
for images, labels, _ in test_loader:
images = Variable(images).to(self.device)
logits1 = self.model1(images)
outputs1 = F.softmax(logits1, dim=1)
_, pred1 = torch.max(outputs1.data, 1)
total1 += labels.size(0)
correct1 += (pred1.cpu() == labels).sum()
correct2 = 0
total2 = 0
for images, labels, _ in test_loader:
images = Variable(images).to(self.device)
logits2 = self.model2(images)
outputs2 = F.softmax(logits2, dim=1)
_, pred2 = torch.max(outputs2.data, 1)
total2 += labels.size(0)
correct2 += (pred2.cpu() == labels).sum()
acc1 = 100 * float(correct1) / float(total1)
acc2 = 100 * float(correct2) / float(total2)
return acc1, acc2
# Train the Model
def train(self, train_loader, epoch):
print('Training ...')
self.model1.train() # Change model to 'train' mode.
self.model2.train() # Change model to 'train' mode
if self.adjust_lr == 1:
self.adjust_learning_rate(self.optimizer, epoch)
train_total = 0
train_correct = 0
train_total2 = 0
train_correct2 = 0
pure_ratio_1_list = []
pure_ratio_2_list = []
for i, (images, labels, indexes) in enumerate(train_loader):
ind = indexes.cpu().numpy().transpose()
if i > self.num_iter_per_epoch:
break
images = Variable(images).to(self.device)
labels = Variable(labels).to(self.device)
# Forward + Backward + Optimize
logits1 = self.model1(images)
prec1 = accuracy(logits1, labels, topk=(1,))
train_total += 1
train_correct += prec1
logits2 = self.model2(images)
prec2 = accuracy(logits2, labels, topk=(1,))
train_total2 += 1
train_correct2 += prec2
loss_1, loss_2, pure_ratio_1, pure_ratio_2 = self.loss_fn(logits1, logits2, labels, self.rate_schedule[epoch],
ind, self.noise_or_not, self.co_lambda)
self.optimizer.zero_grad()
loss_1.backward()
self.optimizer.step()
pure_ratio_1_list.append(100 * pure_ratio_1)
pure_ratio_2_list.append(100 * pure_ratio_2)
if (i + 1) % self.print_freq == 0:
print(
'Epoch [%d/%d], Iter [%d/%d] Training Accuracy1: %.4F, Training Accuracy2: %.4f, Loss1: %.4f, Loss2: %.4f, Pure Ratio1 %.4f %% Pure Ratio2 %.4f %%'
% (epoch + 1, self.n_epoch, i + 1, len(self.train_dataset) // self.batch_size, prec1, prec2,
loss_1.data.item(), loss_2.data.item(), sum(pure_ratio_1_list) / len(pure_ratio_1_list), sum(pure_ratio_2_list) / len(pure_ratio_2_list)))
train_acc1 = float(train_correct) / float(train_total)
train_acc2 = float(train_correct2) / float(train_total2)
return train_acc1, train_acc2, pure_ratio_1_list, pure_ratio_2_list
def adjust_learning_rate(self, optimizer, epoch):
for param_group in optimizer.param_groups:
param_group['lr'] = self.alpha_plan[epoch]
param_group['betas'] = (self.beta1_plan[epoch], 0.999) # Only change beta1