def init_model(model_type):
model = []
if model_type == 'LeNet5':
model = fedmodel.LeNet5()
elif model_type == 'MLP':
model = fedmodel.MLP()
elif model_type == 'ResNet18':
model = fedmodel.ResNet18()
elif model_type == 'CNN':
model = fedmodel.CNN()
return model
def main():
args = parse_args()
train_dataset, test_dataset = dataset.get_dataset(args.path,
args.use_augmentation,
args.use_fivecrop)
train_loader = DataLoader(train_dataset,
args.batch,
True,
num_workers=args.worker,
pin_memory=True)
test_loader = DataLoader(test_dataset,
args.batch,
False,
num_workers=args.worker,
pin_memory=True)
if args.cuda:
torch.cuda.set_device(0)
device = torch.device('cuda')
else:
device = torch.device('cpu')
if args.model == 'ResNet18':
mymodel = model.ResNet18(args.frozen_layers).to(device)
elif args.model == 'ResNet34':
mymodel = model.ResNet34(args.frozen_layers).to(device)
elif args.model == 'ResNet50':
mymodel = model.ResNet50(args.frozen_layers).to(device)
elif args.model == 'DenseNet':
mymodel = model.DenseNet().to(device)
else:
pass
op = optim.Adam(mymodel.parameters(), lr=args.lr)
train_losses, test_mF1s, test_precisions, test_recalls = [], [], [], []
early = args.early
for i in range(args.epoch):
train_loss = train.train(mymodel, op, train_loader, i, device,
args.log, utils.pos_weight)
mF1, recall, presicion = test.test(mymodel, test_loader, device,
args.use_fivecrop)
train_losses.append(train_loss)
test_mF1s.append(mF1)
test_precisions.append(presicion)
test_recalls.append(recall)
early = utils.early_stop(test_mF1s, early)
if early <= 0:
break
utils.save_log(mymodel, train_losses, test_mF1s, test_precisions,
test_recalls)
def main():
args = parse_args()
save = True
use_gpu = args.cuda == 'True'
load = args.load == 'True'
train_tfs = transforms.Compose([
transforms.Resize(299),
transforms.RandomSizedCrop(299),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
ds = my_dataset("train", train_tfs)
dataset_size = ds.__len__()
print(dataset_size)
train_ds, val_ds = torch.utils.data.random_split(ds, [13000, 1463])
train_loader = torch.utils.data.DataLoader(train_ds,
args.BS,
False,
num_workers=8)
val_loader = torch.utils.data.DataLoader(val_ds,
args.BS,
False,
num_workers=8)
print('train: ', len(train_ds))
print('validation:', len(val_ds))
print(type(ds), type(train_ds))
if args.model == 'ResNet18':
test_model = model.ResNet18()
if args.model == 'ResNet50':
test_model = model.ResNet50()
if args.model == 'Inception':
test_model = model.Inception()
if args.model == 'DenseNet':
test_model = model.DenseNet()
if use_gpu:
test_model = test_model.cuda()
if load:
test_model.load_state_dict(torch.load('params' + args.model + '.pkl'))
optimizer = optim.Adam(test_model.parameters(), lr=args.lr)
print(use_gpu)
result = train(test_model, args.epoch, optimizer, train_loader, val_loader,
args.model, save, use_gpu)
test(result, val_loader, use_gpu)
DATASET = dataset.Testset(args.data_path)
DATALOADER = Data.DataLoader(DATASET,
batch_size=args.batch_size,
shuffle=True,
num_workers=20)
NUM_CLASSES = DATASET.num_classes
print('Data path: ' + args.data_path)
print('Number of classes: %d' % NUM_CLASSES)
print('Batch size: %d' % args.batch_size)
print('Epoch size: %d' % args.epoch_size)
num_batches = len(DATALOADER)
batch_total = num_batches * args.epoch_size
if args.model == 'ResNet18':
MODEL = model.ResNet18(pseudo=args.pseudo)
elif args.model == 'ResNet34':
MODEL = model.ResNet34(pseudo=args.pseudo)
elif args.model == 'ResNet50':
MODEL = model.ResNet50(pseudo=args.pseudo)
elif args.model == 'ResNet101':
MODEL = model.ResNet101(pseudo=args.pseudo)
ARCFACE = lossfunction.Arcface(512, NUM_CLASSES)
if args.optim == 'Adam':
OPTIMIZER = torch.optim.Adam([{
'params': MODEL.parameters()
}, {
'params': ARCFACE.parameters()
}],
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import ReduceLROnPlateau
from pre_process import transform_train, transform_test
# whether use gpu
use_cuda = torch.cuda.is_available()
# default parameters
DATA_ROOT = '../data/'
num_epochs = 50
batch_size = 128
model_names = {
'dnn': model.DNN(3072, 4096, 10),
'cnn': model.CNN(),
'resnet18': model.ResNet18(),
'resnet34': model.ResNet34(),
'resnet50': model.ResNet50()
}
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('--model_type',
type=str,
default='dnn',
help="the type of model")
parser.add_argument('--lr',
type=float,
default=0.1,
help='the initial learning rate')
import torchvision.transforms as transforms
import torch.optim as optim
from torchsummary import summary
# Importing Modules
import data as d
import show_images as s
import model as m
import train_test as t
classes, trainloader, testloader = d.load()
s.show_random_images(trainloader, classes)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print(device)
model = m.ResNet18().to(device)
summary(model, input_size=(3, 32, 32))
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=0.1,
momentum=0.9,
weight_decay=0.0005)
testLoss = []
testAcc = []
EPOCHS = 1
for epoch in range(EPOCHS):
print("EPOCH:", epoch)
print("Device:", device)
t.train(model, device, trainloader, optimizer, criterion, epoch)
test_loss, test_acc = t.test(model, device, criterion, testloader)
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])
transforms_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])
trainset = datasets.CIFAR10('../../data', train=True, download=is_download, transform=transforms_train)
testset = datasets.CIFAR10('../../data', train=False, download=is_download, transform=transforms_test)
trainloader = data.DataLoader(trainset, batch_size=size_train, shuffle=True, num_workers=8)
testloader = data.DataLoader(testset, batch_size=size_test, shuffle=False, num_workers=8)
# model
predictor = model.ResNet18()
# net
net = experiment.E_basic(predictor)
if use_cuda:
net.cuda()
updater = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=decay)
trainer = Trainer(use_cuda, trainloader, max_epoch, net, updater)
print('trainer completed')
trainer.headtrain(extensions.basic_load, args.resume_epoch, 'accuracy')
trainer.headepoch(extensions.drop_lr, lr_trigger, lrs)
# set tailepoch extensions, the work to do at the end of each epoch
def main_worker(gpu, ngpus_per_node, time_str, args):
global best_acc1
args.gpu = gpu
prune_rate = float(args.prune_rate)
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
torch.cuda.empty_cache()
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
os.environ['NCCL_DEBUG'] = 'INFO'
os.environ['NCCL_SOCKET_IFNAME'] = 'eno1'
os.environ['NCCL_IB_DISABLE'] = '1'
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = "29595"
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create models
if args.pretrained:
print("=> using pre-trained models '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
elif args.arch == 'ResNet18':
model = self_model.ResNet18(
conv_act='relu',
train_mode_conv='Sign_symmetry_magnitude_uniform',
train_mode_fc='Sign_symmetry_magnitude_uniform',
prune_flag='StochasticFA',
prune_percent=prune_rate,
angle_measurement=False)
elif args.arch == 'AlexNet':
model = self_model.AlexNet_SFA(
train_mode_conv='Sign_symmetry_magnitude_uniform',
train_mode_fc='Sign_symmetry_magnitude_uniform',
prune_rate=prune_rate,
angle_measurement=False)
else:
print("=> creating models '{}'".format(args.arch))
model = models.__dict__[args.arch]()
if not torch.cuda.is_available():
print('using CPU, this will be slow')
elif args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(
(args.workers + ngpus_per_node - 1) / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
if args.gpu is None:
checkpoint = torch.load(args.resume)
else:
# Map models to be loaded to specified single gpu.
loc = 'cuda:{}'.format(args.gpu)
checkpoint = torch.load(args.resume, map_location=loc)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion, args)
return
if args.cos_annl_lr_scheduler:
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=200)
with torch.cuda.device(0):
macs, params = get_model_complexity_info(model, (3, 227, 227),
as_strings=True,
print_per_layer_stat=True,
verbose=True)
logging.info('{:<30} {:<8}'.format('Computational complexity: ',
macs))
logging.info('{:<30} {:<8}'.format('Number of parameters: ', params))
# summary(models, input_size=(3, 227, 227), device='cpu')
logging.info('Model:')
logging.info(model)
layers_zero_grad_df = pd.DataFrame()
zero_grads_percentage_list = []
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
localtime = time.localtime(time.time())
print('fred: Now is {}d-{}h-{}m-{}s'.format(localtime.tm_mday,
localtime.tm_hour,
localtime.tm_min,
localtime.tm_sec))
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch,
layers_zero_grad_df, zero_grads_percentage_list, time_str, args)
# step the scheduler
if args.cos_annl_lr_scheduler:
scheduler.step()
# evaluate on validation set
acc1 = validate(val_loader, model, criterion, args)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if not args.multiprocessing_distributed or (
args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1':
best_acc1, # the epoch might not be best epoch, with the best_acc1 stored, the model_best
# pth.tar won't be updated unless it is the real best in this trial
'optimizer': optimizer.state_dict(),
},
is_best,
arch=args.arch,
foldername=args.folder_name)
parser.add_argument(
'--outf',
default='./checkpoint',
help='folder to output images and model checkpoints') #输出结果保存路径
args = parser.parse_args()
# 定义是否使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 超参数设置
EPOCH = 2 # 遍历数据集次数
pre_epoch = 0 # 定义已经遍历数据集的次数
LR = 0.1 # 学习率
# 模型定义-ResNet
net = model.ResNet18().to(device)
# 定义损失函数和优化方式
criterion = nn.CrossEntropyLoss() # 损失函数为交叉熵，多用于多分类问题
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9, weight_decay=5e-4)
# 优化方式为mini-batch momentum-SGD，并采用L2正则化（权重衰减）
# 训练
if __name__ == "__main__":
print("Start Training, Resnet-18!")
num_iters = 0
for epoch in range(pre_epoch, EPOCH):
print('\nEpoch: %d' % (epoch + 1))
net.train()
sum_loss = 0.0
correct = 0.0
