- https://github.com/fossasia/visdom
- Step 1: 导入工具包
from visdom import Visdom # ⭐️
- Step 2: 模型定义中初始化三个函数 （
self.global_step = 0记录全局次数）
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.global_step = 0
...
def vis_init(self): # ⭐️
...
def vis_moni_train(self,loss,acc,global_step): # ⭐️
...
def vis_moni_test(self,loss,acc,global_step): # ⭐️
...
- Step 3: 调用Visdom初始化
model = CNN()
model.vis_init() if params.visdom else None # ⭐️
- Step 4:
for epoch in Epoch:
for batch in Batch:
...
losstr =
correct =
model.vis_moni_train(losstr.item(), correct, model.global_step) # ⭐️
model.vis_moni_test(loss,correct,epoch) # ⭐️
- 案例代码
import numpy as np
import matplotlib.pyplot as plt
import torch
import scipy.io as sio
from torch.utils.data import Dataset, TensorDataset, DataLoader
import torch
import torchvision
from torchvision import transforms
from torch import nn
from torch import optim
from torch.autograd import Variable
import torch.nn.functional as F
from visdom import Visdom # ⭐️
import time
import argparse
import sys
import os
import pprint
# from utils.func import TaskScheduler
# from utils.visulization import vis_confusion_matrix_all
import sys
import os
# sys.path.append('D:\ImbalanceRevise')
sys.path.append(os.getcwd())
from utils import task_pre, get_save_name
# import utils.task_pre as task_pre
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
import warnings
warnings.filterwarnings('ignore')
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.global_step = 0
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout1 = nn.Dropout2d(0.25)
self.dropout2 = nn.Dropout2d(0.5)
self.fc1 = nn.Linear(12544, 128)
self.fc2 = nn.Linear(128, 10)
self.bat_acc_list = []
self.bat_los_list = []
self.fit_acc_list = []
self.fit_los_list = []
self.val_acc_list = []
self.val_los_list = []
self.pre_acc_list = []
self.pre_los_list = []
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
output = F.log_softmax(x, dim=1)
return output
def vis_init(self): # ⭐️
self.viz = Visdom()
self.viz.line([0.], [0.], win='train_loss', opts=dict(title='train loss'))
self.viz.line([0.], [0.], win='train_acc', opts=dict(title='train acc'))
self.viz.line([0.], [0.], win='val_loss', opts=dict(title='val loss'))
self.viz.line([0.], [0.], win='val_acc', opts=dict(title='val acc'))
def vis_moni_train(self,loss,acc,global_step): # ⭐️
self.viz.line([loss], [global_step], win='train_loss', update='append')
self.viz.line([acc], [global_step], win='train_acc', update='append')
def vis_moni_test(self,loss,acc,global_step): # ⭐️
self.viz.line([loss], [global_step], win='val_loss', update='append')
self.viz.line([acc], [global_step], win='val_acc', update='append')
def fit(self, params, model, device, train_loader, epoch, optimizer):
# print('Epochs: ', epoch)
correct = 0
L = len(train_loader.dataset)
model.train()
for batch_idx, (data, target) in enumerate(train_loader): # Load from cpu
# print(batch_idx)
data, target = data.to(device), target.to(device) # Data to device
optimizer.zero_grad() # Zero gradient
output = model(data) # Forward propagation
losstr = F.nll_loss(output,target) # Calculate loss
losstr.backward() # Back propagation
optimizer.step() # Optimize parameters
predict = output.argmax(dim=1, keepdim=True) # Get the index
correct = predict.eq(target.view_as(predict)).sum().item()/len(target) # Sum up batch acc
self.bat_acc_list.append(correct)
self.bat_los_list.append(losstr.cpu().detach().numpy())
model.global_step += 1
if params.visdom:
model.vis_moni_train(losstr.item(), correct, model.global_step) # ⭐️
def record(self, params, model, device, data_loader, epoch, phrase):
# * Get validation accuracy per epoch
model.eval（)
loss = 0
correct = 0
with torch.no_grad(): # Save memory
for data, target in data_loader: # Load from cpu
data, target = data.to(device), target.to(device) # Data to device
output = model(data) # Forward propagation
loss += F.nll_loss(output, target, reduction='sum').item() # Sum up batch loss
predict = output.argmax(dim=1, keepdim=True) # Get the index
correct += predict.eq(target.view_as(predict)).sum().item() # Sum up batch acc
# Loss, accuracy and visualize
loss /= len(data_loader.dataset)
correct /= len(data_loader.dataset)
if params.visdom and phrase=='Valida':
model.vis_moni_test(loss,correct,epoch) # ⭐️
model.test_acc = correct
model.test_pred = predict.view(-1).cpu().numpy().reshape(-1,1)
try: # cpu
target, predict = target.numpy(), predict.numpy().ravel()
except: # gpu
target, predict = target.cpu().numpy(), predict.cpu().numpy().ravel()
acc = accuracy_score(target, predict)
if phrase=='Source':
self.fit_acc_list.append(acc)
self.fit_los_list.append(loss)
if phrase=='Valida':
self.val_acc_list.append(acc)
self.val_los_list.append(loss)
if phrase=='Target':
self.pre_acc_list.append(acc)
self.pre_los_list.append(loss)
return acc, loss
def train_CNN(params, Xs, Ys, Xt, Yt):
# * [1] Load data
Xs = Xs.reshape(-1,1,32,32)
Xt = Xt.reshape(-1,1,32,32)
# * [2] Set device
if params.gpuidx is None:
device = 'cpu'
if isinstance(params.gpuidx, int) and torch.cuda.is_available():
device = 'cuda:{}'.format(params.gpuidx)
print('Device: {}'.format(device))
print('-'*80)
# params.batch = 32
# params.Epoch = 50
params.model = 'CNN'
# * [4] Dataloader
Xs, Xt = torch.Tensor(Xs), torch.Tensor(Xt)
Ys, Yt = torch.LongTensor(Ys), torch.LongTensor(Yt)
# Xs, Xv, Ys, Yv = train_test_split(Xs, Ys, test_size=0.3)
_, Xv, _, Yv = train_test_split(Xs, Ys, test_size=0.3)
Xs, Xv, Ys, Yv = train_test_split(Xs, Ys, test_size=0.1)
source_dataset = TensorDataset(torch.Tensor(Xs), torch.LongTensor(Ys))
valida_dataset = TensorDataset(torch.Tensor(Xv), torch.LongTensor(Yv))
target_dataset = TensorDataset(torch.Tensor(Xt), torch.LongTensor(Yt))
Xs_loader = DataLoader(source_dataset,batch_size=params.batchs,shuffle=True, pin_memory=True)
Xv_loader = DataLoader(valida_dataset,batch_size=len(Xv),shuffle=False, pin_memory=True)
Xt_loader = DataLoader(target_dataset,batch_size=len(Xt),shuffle=False, pin_memory=True)
# * [5] Set optimizer
model = CNN()
model.to(device)
# if params.visdom:
# import subprocess
# command1 = subprocess.Popen(os.system('python -m visdom.server'))
# command1 = subprocess.run(os.system('python -m visdom.server'), check=False)
# os.system('python -m visdom.server')
model.vis_init() if params.visdom else None # ⭐️
optimizer = optim.Adam(model.parameters(),lr=0.001)
# * [6] Start training
source_time = []
target_time = []
for epoch in range(1, params.epochs+1):
t1 = time.time() # Train start
model.fit(params, model, device, Xs_loader, epoch, optimizer)
t2 = time.time() # Train end and Test start
fit_acc, fit_loss = model.record(params, model, device, Xs_loader, epoch, 'Source')
val_acc, val_loss = model.record(params, model, device, Xv_loader, epoch, 'Valida')
pre_acc, pre_loss = model.record(params, model, device, Xt_loader, epoch, 'Target')
print('[Epoch {:0>2d}] Fit acc: {:.3f} loss {:.3f}, Val acc: {:.3f} loss {:.3f}, Pre acc: {:.3f} loss {:.3f}.'.format(epoch, fit_acc, fit_loss, val_acc, val_loss, pre_acc, pre_loss))
t3 = time.time() # Test end
source_time.append(t2-t1)
target_time.append(t3-t2)
source_time = sum(source_time)
# model.tagret_time = sum(target_time)/params.Epoch
# print('Training time for {} on device [{}] is {:.3f}s'.format(params.model, device, model.source_time))
mdict={ 'bat_acc': model.bat_acc_list,
'bat_los': model.bat_los_list,
'fit_acc': model.fit_acc_list,
'fit_los': model.fit_los_list,
'val_acc': model.val_acc_list,
'val_los': model.val_los_list,
'pre_acc': model.pre_acc_list,
'pre_los': model.pre_los_list,
'Acc': pre_acc,
'Tim': source_time}
save_name = get_save_name(params)
print(save_name)
sio.savemat(save_name, mdict=mdict)
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", type=str, default='PHM')
parser.add_argument("--datafea", type=str, default='RMS')
parser.add_argument("--source", type=int, default=3)
parser.add_argument("--target", type=int, default=0)
parser.add_argument("--method", type=str, default='CNN')
parser.add_argument("--iratio", type=float, default=1)
parser.add_argument("--repeat", type=int, default=5)
parser.add_argument("--remark", type=str, default='')
parser.add_argument("--folder", type=str, default='Exp1_balance')
parser.add_argument("--gpuidx", type=int, default=0)
parser.add_argument("--visdom", action='store_true')
parser.add_argument("--batchs", type=int, default=64)
parser.add_argument("--epochs", type=int, default=30)
params, unknown = parser.parse_known_args()
print('-'*80)
pprint.pprint(vars(params))
print('-'*80)
Xs, Ys, Xt, Yt = task_pre(params)
train_CNN(params, Xs, Ys, Xt, Yt)
