def main(args):
test_ds = MnistDataset(
args.test_image_file,
args.test_label_file,
transform=transforms.Compose([ToTensor()]),
)
test_loader = torch.utils.data.DataLoader(
test_ds,
batch_size=args.batch_size,
collate_fn=collate_fn,
shuffle=False,
)
model = Net().to(device)
model.load_state_dict(torch.load(args.checkpoint))
model.eval()
predicts = []
truths = []
with torch.no_grad():
for i, sample in enumerate(test_loader):
X, Y_true = sample["X"].to(device), sample["Y"].to(device)
output = model(X)
predicts.append(torch.argmax(output, dim=1))
truths.append(Y_true)
predicts = torch.cat(predicts, dim=0)
truths = torch.cat(truths, dim=0)
acc = torch.sum(torch.eq(predicts, truths))
print("Acc: {:.4f}".format(acc / len(predicts)))
else:
checkpoint = torch.load('./weights/best_model-20200904.pth.tar')
#new_state_dict = OrderedDict()
# 用了nn.DataParallel的模型需要处理才能在cpu上使用
'''for k, v in checkpoint.items():
name = k[7:] # remove module.
new_state_dict[name] = v
model.load_state_dict(new_state_dict)'''
#model.load_state_dict(torch.load('./weights/best-8-24.pth.tar'))
model.load_state_dict(torch.load('./weights/best_model-20200904.pth.tar'))
model.eval()
model = model.to(device)
a = 0
b = 0
start_time = time.time()
for file in tqdm(os.listdir(data_path)):
img_path = os.path.join(data_path, file)
img = cv2.imread(img_path)
#img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
img = Image.fromarray(img)
frame = transform(img)
train_correct = (train_pred == targets).sum()
train_acc += train_correct.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
'''#print statisics
running_loss+=loss.item()
if epoch%10==9:
print(print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 10)))
running_loss=0.0'''
net.eval()
eval_loss = 0.
eval_acc = 0.
for inputs, targets in dataloaders['valid']:
inputs = inputs.to(device)
targets = targets.to(device)
predictions = net(inputs)
loss = criterion(predictions, targets)
eval_loss += loss.item()
eval_pred = torch.max(predictions, 1)[1]
num_correct = (eval_pred == targets).sum()
eval_acc += num_correct.item()
class Trainer(object):
def __init__(self, args):
self.args = args
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.prepare_data()
self.setup_train()
def prepare_data(self):
train_val = MnistDataset(
self.args.train_image_file,
self.args.train_label_file,
transform=transforms.Compose([ToTensor()]),
)
train_len = int(0.8 * len(train_val))
train_ds, val_ds = torch.utils.data.random_split(
train_val, [train_len, len(train_val) - train_len]
)
print("Train {}, val {}".format(len(train_ds), len(val_ds)))
self.train_loader = torch.utils.data.DataLoader(
train_ds,
batch_size=self.args.batch_size,
collate_fn=collate_fn,
shuffle=True,
)
self.val_loader = torch.utils.data.DataLoader(
val_ds,
batch_size=self.args.batch_size,
collate_fn=collate_fn,
shuffle=False,
)
def setup_train(self):
self.model = Net().to(self.device)
self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.args.lr)
self.criterion = nn.CrossEntropyLoss().to(self.device)
if not os.path.isdir(self.args.ckpt):
os.mkdir(self.args.ckpt)
def train_one_epoch(self):
train_loss = 0.0
self.model.train()
for i, sample in enumerate(self.train_loader):
X, Y_true = sample["X"].to(self.device), sample["Y"].to(self.device)
self.optimizer.zero_grad()
output = self.model(X)
loss = self.criterion(output, Y_true)
loss.backward()
self.optimizer.step()
train_loss += loss.item()
return train_loss / len(self.train_loader)
def evaluate(self):
val_loss = 0.0
self.model.eval()
predicts = []
truths = []
with torch.no_grad():
for i, sample in enumerate(self.val_loader):
X, Y_true = sample["X"].to(self.device), sample["Y"].to(self.device)
output = self.model(X)
loss = self.criterion(output, Y_true)
val_loss += loss.item()
predicts.append(torch.argmax(output, dim=1))
truths.append(Y_true)
predicts = torch.cat(predicts, dim=0)
truths = torch.cat(truths, dim=0)
acc = torch.sum(torch.eq(predicts, truths))
return acc / len(predicts), val_loss / (len(self.val_loader))
def run(self):
min_loss = 10e4
max_acc = 0
for epoch in range(self.args.epochs):
train_loss = self.train_one_epoch()
val_acc, val_loss = self.evaluate()
if val_acc > max_acc:
max_acc = val_acc
torch.save(
self.model.state_dict(),
os.path.join(
self.args.ckpt,
"{}_{}_{:.4f}.pth".format(self.args.name, epoch, max_acc),
),
)
print(
"Epoch {}, loss {:.4f}, val_acc {:.4f}".format(
epoch, train_loss, val_acc
)
)