def train(args):
my_dataset = MyDataset("../data/train", transform=x_transforms, target_transform=y_transforms)
dataloaders = DataLoader(my_dataset, batch_size=args.batch_size, shuffle=True, num_workers=1)
model = Unet(3, 1).to(device)
model.train()
criterion = torch.nn.BCELoss()
optimizer = optim.Adam(model.parameters())
num_epochs = args.epochs
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
data_size = len(dataloaders.dataset)
epoch_loss = 0
step = 0
for x, y in dataloaders:
step += 1
inputs = x.to(device)
lables = y.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, lables)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
print("%d/%d, train_loss:%0.3f" % (step, (data_size - 1) // dataloaders.batch_size + 1, loss.item()))
print("epoch %d loss:%0.3f" % (epoch, epoch_loss))
torch.save(model.state_dict(), 'model_weights.pth')
return model
def train():
save_dir = "/home/FuDawei/NLP/SQUAD/unet/data/"
train_examples, dev_examples, opt = prepare_train(save_dir)
epoch = 30
batch_size = 32
model = Unet(opt=opt).to(device)
parameters = filter(lambda p: p.requires_grad, model.parameters())
optimizer = torch.optim.Adamax(parameters, lr=opt["lr"])
best_score, exact_scores, f1_scores = 0, [], []
count = 0
total_loss = 0
for ep in range(epoch):
model.train()
for batch_data in get_batch_data(train_examples, batch_size):
data = model.get_data(batch_data)
loss = model(data)
model.zero_grad()
optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(parameters, 10)
optimizer.step()
model.reset_parameters()
count += 1
# print(loss.item())
# Evaluate(dev_examples, model)
total_loss += loss.item()
if count % 100 == 0:
print(total_loss / 100)
total_loss = 0
# model.eval()
# Evaluate(dev_examples, model, opt)
if not opt["fix_word_embedding"]:
model.reset_parameters()
print(ep)
model.eval()
exact, f1 = Evaluate(dev_examples, model, opt)
exact_scores.append(exact)
f1_scores.append(f1)
if f1 > best_score:
best_score = f1
torch.save(model.state_dict(), save_dir + "best_model")
with open(save_dir + '_f1_scores.pkl', 'wb') as f:
pkl.dump(f1_scores, f)
with open(save_dir + '_exact_scores.pkl', 'wb') as f:
pkl.dump(exact_scores, f)
def main(FLAGS):
"train and validate the Unet model"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#data directory
data_dir = FLAGS.dataset_dir
#log_directory
log_dir = FLAGS.log_dir
# Hyper and other parameters
train_batch_size = FLAGS.train_batch_size
val_batch_size = FLAGS.val_batch_size
aug_flag = FLAGS.aug
num_epochs = FLAGS.epochs
num_classes = 2
# get the train and validation dataloaders
dataloaders = get_dataloaders(data_dir, train_batch_size, val_batch_size,
aug_flag)
model = Unet(3, num_classes)
# Uncomment to run traiing on Multiple GPUs
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model, device_ids=[0, 1])
else:
print("no multiple gpu found")
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=0.02,
momentum=0.9,
weight_decay=0.0005)
#optimizer = optim.Adam(model.parameters(),lr = learning_rate)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
plotter = VisdomLinePlotter(env_name='Unet Train')
# uncomment for leraning rate schgeduler..
train_val(dataloaders, model, criterion, optimizer, num_epochs, log_dir,
device)
train_set, batch_size=args.batch_size, shuffle=True, num_workers=8)
dataloaders['val'] = torch.utils.data.DataLoader(val_set,
batch_size=1,
shuffle=True,
num_workers=8)
from model import Unet
model = Unet(args.in_channels, args.out_channels)
if args.load_caffe_unet:
caffe_unet_path = '/mnt/ccvl15/yixiao/kaggle/models/pretrained/unet.pt'
unet_caffemodel_weights = torch.load(caffe_unet_path)
model = load_conv_weights(model, unet_caffemodel_weights)
elif args.load != '':
model.load_state_dict(torch.load(args.load))
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9)
from losses import cross_entropy_with_soft_dice_2
# loss_func = torch.nn.CrossEntropyLoss()
from criteria import dice, multichannel_dice
try:
model, val_history = train_model(model, dataloaders,
cross_entropy_with_soft_dice_2,
multichannel_dice, optimizer,
args.epochs)
except KeyboardInterrupt:
sys.exit(0)
def train_net(image_size, batch_size, num_epochs, lr, num_workers, checkpoint):
train_loader, val_loader = data_loaders(image_size=(image_size,
image_size),
batch_size=batch_size)
device = torch.device('cuda') if torch.cuda.is_available() else 'cpu'
model = Unet().to(device)
if checkpoint:
model.load_state_dict(torch.load(checkpoint))
criterion = DiceLoss().to(device)
optimizer = Adam(model.parameters(), lr=lr)
logging.info(f'Start training:\n'
f'Num epochs: {num_epochs}\n'
f'Batch size: {batch_size}\n'
f'Learning rate: {lr}\n'
f'Num workers: {num_workers}\n'
f'Scale image size: {image_size}\n'
f'Device: {device}\n'
f'Checkpoint: {checkpoint}\n')
train_losses = []
val_losses = []
for epoch in range(num_epochs):
print(f'Epoch {epoch+1}: ')
train_batch_losses = []
val_batch_losses = []
best_val_loss = 9999
for x_train, y_train in tqdm(train_loader):
x_train = x_train.to(device)
y_train = y_train.to(device)
y_pred = model(x_train)
optimizer.zero_grad()
loss = criterion(y_pred, y_train)
train_batch_losses.append(loss.item())
loss.backward()
optimizer.step()
train_losses.append(sum(train_batch_losses) / len(train_batch_losses))
print(
f'-----------------------Train loss: {train_losses[-1]} -------------------------------'
)
for x_val, y_val in tqdm(val_loader):
x_val = x_val.to(device)
y_val = y_val.to(device)
y_pred = model(x_val)
loss = criterion(y_pred, y_val)
val_batch_losses.append(loss.item())
val_losses.append(sum(val_batch_losses) / len(val_batch_losses))
print(
f'-----------------------Val loss: {val_losses[-1]} -------------------------------'
)
if val_losses[-1] < best_val_loss:
best_val_loss = val_losses[-1]
if not os.path.isdir('weights/'):
os.mkdir('weights/')
torch.save(model.state_dict(), f'weights/checkpoint{epoch+1}.pth')
print(f'Save checkpoint in: weights/checkpoint{epoch+1}.pth')
def train():
model = Unet(input_channel=opt.input_channel, cls_num=opt.cls_num)
model_name = 'Unet_bn'
train_logger = LogWay(
datetime.datetime.now().strftime('%Y-%m-%d %H-%M-%S') + '.txt')
train_data = My_Dataset(opt.train_images, opt.train_masks)
train_dataloader = DataLoader(train_data,
batch_size=opt.batch_size,
shuffle=True,
num_workers=0)
if opt.cls_num == 1:
criterion = torch.nn.BCELoss()
else:
criterion = torch.nn.NLLLoss()
if use_gpu:
model.cuda()
if opt.cls_num == 1:
criterion = torch.nn.BCELoss().cuda()
else:
criterion = torch.nn.NLLLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
lr=opt.learning_rate,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
for epoch in range(opt.epoch):
loss_sum = 0
for i, (data, target) in enumerate(train_dataloader):
data = Variable(data)
target = Variable(target)
if use_gpu:
data = data.cuda()
target = target.cuda()
outputs = model(data)
if opt.cls_num == 1:
outputs = F.sigmoid(outputs).view(-1)
mask_true = target.view(-1)
loss = criterion(outputs, mask_true)
else:
outputs = F.LogSoftmax(outputs, dim=1)
loss = criterion(outputs, target)
loss_sum = loss_sum + loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("epoch:{} batch:{} loss:{}".format(epoch + 1, i,
loss.item()))
info = 'Time:{} Epoch:{} Loss_avg:{}\n'.format(
str(datetime.datetime.now()), epoch + 1, loss_sum / (i + 1))
train_logger.add(info)
adjusting_rate(optimizer, opt.learning_rate, epoch + 1)
realepoch = epoch + 1
if (realepoch % 10 == 0):
save_name = datetime.datetime.now().strftime(
'%Y-%m-%d %H-%M-%S') + ' ' + model_name + str(
realepoch) + '.pt'
torch.save(model.state_dict(), save_name)
def unet_train():
batch_size = 1
num_epochs = [5000, 5000, 5000, 5000, 5000, 5000, 5000, 5000]
num_workers = 2
lr = 0.0001
losslist = ['dice'] # ['focal', 'bce', 'dice', 'lovasz']
optimlist = ['adam'] # ['adam', 'sgd']
iflog = True
SC_root_dir = '../dataset-EdmSealedCrack-512'
train_files, val_files, test_files = myutils.organize_SC_files(SC_root_dir)
train_RC_dataset = DatasetRealCrack('../dataset-EdmCrack600-512/A/train',
transform=transform)
train_SC_dataset = DatasetSealedCrack(files=train_files,
root_dir=SC_root_dir,
transform=data_Train_transforms)
val_RC_dataset = DatasetRealCrack('../dataset-EdmCrack600-512/A/val',
transform=transform)
val_SC_dataset = DatasetSealedCrack(files=val_files,
root_dir=SC_root_dir,
transform=data_Test_transforms)
train_loader = torch.utils.data.DataLoader(ConcatDataset(
train_RC_dataset, train_SC_dataset),
batch_size=2,
shuffle=True,
num_workers=2)
criterion = nn.BCELoss()
focallos = FocalLoss(gamma=2)
doubleFocalloss = focalloss.FocalLoss_2_datasets(gamma=2)
epoidx = -1
for los in losslist:
for opt in optimlist:
start = time.time()
print(los, opt)
torch.manual_seed(77)
torch.cuda.manual_seed(77)
#################
#unet = Unet_SpatialPyramidPooling(3).cuda()
#################
unet = Unet(3).cuda()
SC_classifier = classifier(64, 2).cuda()
RC_classifier = classifier(64, 2).cuda()
##################
#unet = smp.Unet('resnet34', encoder_weights='imagenet').cuda()
#unet.segmentation_head = torch.nn.Sequential().cuda()
#SC_classifier = classifier(16, 2).cuda()
#RC_classifier = classifier(16, 2).cuda()
#UNCOMMENT TO KEEP TRAINING THE BEST MODEL
prev_epoch = 0 # if loading model 58, change to prev_epoch = 58. When saving the model, it is going to be named as 59, 60, 61...
#unet.load_state_dict(torch.load('trained_models/unet_adam_dice_58.pkl'))
#SC_classifier.load_state_dict(torch.load('trained_models/SC_classifier_adam_dice_58.pkl'))
#RC_classifier.load_state_dict(torch.load('trained_models/RC_classifier_adam_dice_58.pkl'))
history = []
if 'adam' in opt:
optimizer = torch.optim.Adam(unet.parameters(), lr=lr)
elif 'sgd' in opt:
optimizer = torch.optim.SGD(unet.parameters(),
lr=10 * lr,
momentum=0.9)
logging.basicConfig(filename='./logs/logger_unet.log',
level=logging.INFO)
total_step = len(train_loader)
epoidx += 1
for epoch in range(num_epochs[epoidx]):
totalloss = 0
for i, (realCrack_batch,
sealedCrack_batch) in enumerate(train_loader):
SC_images = sealedCrack_batch[0].cuda()
SC_masks = sealedCrack_batch[1].cuda()
RC_images = realCrack_batch[0].cuda()
RC_masks = realCrack_batch[1].cuda()
SC_encoder = unet(SC_images)
RC_encoder = unet(RC_images)
#############
SC_outputs = SC_classifier(SC_encoder)
RC_outputs = RC_classifier(RC_encoder)
#############
#Deep lab v3
#SC_outputs = SC_classifier(SC_encoder['out'])
#RC_outputs = RC_classifier(RC_encoder['out'])
##############
if 'bce' in los:
masks = onehot(masks)
loss = criterion(outputs, masks)
elif 'dice' in los:
branch_RC = {'outputs': RC_outputs, 'masks': RC_masks}
branch_SC = {'outputs': SC_outputs, 'masks': SC_masks}
loss = dice_loss_2_datasets(branch_RC, branch_SC)
#masks = onehot(masks)
#loss = dice_loss(outputs, masks)
elif 'lovasz' in los:
masks = onehot(masks)
loss = L.lovasz_hinge(outputs, masks)
elif 'focal' in los:
#loss = focallos(outputs, masks.long())
branch_RC = {
'outputs': RC_outputs,
'masks': RC_masks.long()
}
branch_SC = {
'outputs': SC_outputs,
'masks': SC_masks.long()
}
loss = doubleFocalloss(branch_RC, branch_SC)
totalloss += loss * RC_images.size(0) #*2?
#print(RC_images.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 10 == 0:
print(epoch, i)
print("total loss: ", totalloss)
if i % 1000 == 0:
print("Epoch:%d; Iteration:%d; Loss:%f" %
(epoch, i, loss))
if i + 1 == total_step: # and epoch%1==0: #and val_miou>0.85:
torch.save(
unet.state_dict(),
'./trained_models/unet_' + opt + '_' + los + '_' +
str(epoch + 1 + prev_epoch) + '.pkl')
torch.save(
RC_classifier.state_dict(),
'./trained_models/RC_classifier_' + opt + '_' +
los + '_' + str(epoch + 1 + prev_epoch) + '.pkl')
torch.save(
SC_classifier.state_dict(),
'./trained_models/SC_classifier_' + opt + '_' +
los + '_' + str(epoch + 1 + prev_epoch) + '.pkl')
history_np = np.array(history)
np.save('./logs/unet_' + opt + '_' + los + '.npy', history_np)
end = time.time()
print((end - start) / 60)
num_workers=4)
generator = define_G(
4,
1,
64,
'unet_128',
norm='instance',
)
discriminator = netD()
unet = Unet()
unet.load_state_dict(torch.load("./weight/unet_pretrained.pth"))
optimizer_g = torch.optim.Adam(generator.parameters(), lr=0.0002)
optimizer_d = torch.optim.Adam(discriminator.parameters(), lr=0.0002)
optimizer_s = torch.optim.Adam(unet.parameters(), lr=0.0002)
generator.cuda()
discriminator.cuda()
unet.cuda()
EPOCH = 100
num_iter = len(train_loader)
D_LOSS = []
G_LOSS = []
# S_LOSS=[]
f = open("./loss_gan.txt", 'a')
print(time.strftime('|---------%Y-%m-%d %H:%M:%S---------|',
time.localtime(time.time())),
file=f)
discriminator.train()
unet.train()
shuffle=True)
num_data = len(dataset_train)
num_epoch = np.ceil(num_data / args.batch_size)
print(f'Number of data: {num_data}')
print(f'Number of epochs: {num_epoch}')
# Model load.
u_net = Unet().to(device)
# loss
bce_loss = nn.BCEWithLogitsLoss().to(device)
# Optimizer.
optim = torch.optim.Adam(u_net.parameters(), lr=args.lr)
# 기타 function 설정
to_numpy = lambda x: x.to('cpu').detach().numpy().transpose(
0, 2, 3, 1) # device 위에 올라간 텐서를 detach 한 뒤 numpy로 변환
denorm = lambda x, mean, std: (x * std) + mean
classifier = lambda x: 1.0 * (
x > 0.5) # threshold 0.5 기준으로 indicator function으로 classifier 구현
# Tensorboard
tensorboard_train = SummaryWriter(
log_dir=os.path.join(args.log_dir, 'train'))
# Training loop
init_epoch = 0
max_num_iters = num_data // args.batch_size
valid_set = Cloudset(
r.train, 'valid', r.valid_ids, r.train_fold,
validation_augmentation_kaggle()
# validation_augmentation()
# try different augmentation here
)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
valid_loader = DataLoader(valid_set, batch_size=batch_size, shuffle=False)
print('training data loaded')
net = Unet(c_in=3, c_out=4).float()
if is_gpu:
net.cuda()
print('unet built')
# training
criterion = BceDiceLoss(eps=1e-1) # make sure tp=1 at least
optimizer = optim.Adam(net.parameters(), lr=initial_lr)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.2,
patience=2,
cooldown=2)
valid_loss_min = np.Inf
# for plot
train_loss_list = []
valid_loss_list = []
lr_list = []
# start
for epoch in range(1, max_epochs + 1):
train_loss = 0.0
valid_loss = 0.0
# train
net.train()
ToTensorV2()
])
train_dataset = KvasirSegDataset(train_path, t_transforms)
val_dataset = KvasirSegDataset(val_path, v_transforms)
train_dataloader = DataLoader(train_dataset, batch_size)
val_dataloader = DataLoader(val_dataset, batch_size)
in_classes = 3
out_classes = 1
model = Unet(in_classes, out_classes)
criteria = DICELoss()
metrics = {'iou': iou}
optim = torch.optim.SGD(model.parameters(),
lr=0.001,
momentum=0.99,
weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.CyclicLR(optim,
base_lr=0.001,
max_lr=0.01)
lr_s = torch.optim.lr_scheduler.CosineAnnealingLR(optim, 100, 0.001)
# train(train_dataloader, val_dataloader, model, epochs, criteria, metrics, optim, scheduler=lr_s, device=device)
predict(
model,
'D:\\Kvasir-SEG',
device,
f_name='unet_aug4_pred_masks',
model_path=
'C:\\Users\\DSKIM\\Google 드라이브\\AI\\medical-projects\\Kvasir-Seg\\unet_aug4_models\\Unet_199_22.pth'
def main(args):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
### Hyperparameters Setting ###
epochs = args.epochs
batch_size = args.batch_size
num_workers = args.num_workers
valid_ratio = args.valid_ratio
threshold = args.threshold
separable = args.separable
down_method = args.down_method
up_method = args.up_method
### DataLoader ###
dataset = DataSetWrapper(batch_size, num_workers, valid_ratio)
train_dl, valid_dl = dataset.get_data_loaders(train=True)
### Model: U-Net ###
model = Unet(input_dim=1,
separable=separable,
down_method=down_method,
up_method=up_method)
model.summary()
model = nn.DataParallel(model).to(device)
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=len(train_dl),
eta_min=0,
last_epoch=-1)
criterion = nn.BCEWithLogitsLoss()
train_losses = []
val_losses = []
###Train & Validation start ###
mIOU_list = []
best_mIOU = 0.
step = 0
for epoch in range(epochs):
### train ###
pbar = tqdm(train_dl)
model.train()
losses = []
for (img, label) in pbar:
optimizer.zero_grad()
img, label = img.to(device), label.to(device)
pred = model(img)
# pred = Padding()(pred, label.size(3))
loss = criterion(pred, label)
loss.backward()
optimizer.step()
losses.append(loss.item())
pbar.set_description(
f'E: {epoch + 1} | L: {loss.item():.4f} | lr: {scheduler.get_lr()[0]:.7f}'
)
scheduler.step()
if (epoch + 1) % 10:
losses = sum(losses) / len(losses)
train_losses.append(losses)
### validation ###
with torch.no_grad():
model.eval()
mIOU = []
losses = []
pbar = tqdm(valid_dl)
for (img, label) in pbar:
img, label = img.to(device), label.to(device)
pred = model(img)
loss = criterion(pred, label)
mIOU.append(get_IOU(pred, label, threshold=threshold))
losses.append(loss.item())
mIOU = sum(mIOU) / len(mIOU)
mIOU_list.append(mIOU)
if (epoch + 1) % 10:
losses = sum(losses) / len(losses)
val_losses.append(losses)
print(
f'VL: {loss.item():.4f} | mIOU: {100 * mIOU:.1f}% | best mIOU: {100 * best_mIOU:.1f}'
)
### Early Stopping ###
if mIOU > best_mIOU:
best_mIOU = mIOU
save_state = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'train_losses': train_losses,
'val_losses': val_losses,
'best_mIOU': best_mIOU
}
torch.save(
save_state,
f'./checkpoint/{down_method}_{up_method}_{separable}.ckpt')
step = 0
else:
step += 1
if step > args.patience:
print('Early stopped...')
return
def train(args, x_train, y_train, x_valid, y_valid):
writer = SummaryWriter()
best_dice = 0
model = Unet().to(args.device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
bce_loss = torch.nn.BCELoss()
train_dataloader = load_dataset(x_train, y_train, args.batch_size, True)
valid_dataloader = load_dataset(x_valid, y_valid, args.batch_size, False)
result = {}
result['train/BCE'] = []
result['train/Dice'] = []
result['valid/BCE'] = []
result['valid/Dice'] = []
for epoch in range(args.epochs):
print('train step: epoch {}'.format(str(epoch+1).zfill(4)))
train_bce = []
train_dice = []
for inp_im, lab_im in tqdm(train_dataloader):
inp_im = inp_im.to(args.device)
lab_im = lab_im.to(args.device)
pred = model(inp_im)
bce = bce_loss(pred, lab_im)
dice = calc_dice(pred, lab_im)
train_bce.append(bce.item())
train_dice.append(dice)
model.zero_grad()
bce.backward()
optimizer.step()
result['train/BCE'].append(statistics.mean(train_bce))
result['train/Dice'].append(statistics.mean(train_dice))
writer.add_scalar('train/BinaryCrossEntropy', result['train/BCE'][-1], epoch+1)
writer.add_scalar('train/DiceScore', result['train/Dice'][-1], epoch+1)
print('BCE: {}, Dice: {}'.format(result['train/BCE'][-1], result['train/Dice'][-1]))
if (epoch+1) % 10 == 0 or (epoch+1) == 1:
with torch.no_grad():
print('valid step: epoch {}'.format(str(epoch+1).zfill(4)))
model.eval()
valid_bce = []
valid_dice = []
for inp_im, lab_im in tqdm(valid_dataloader):
inp_im = inp_im.to(args.device)
lab_im = lab_im.to(args.device)
pred = model(inp_im)
bce = bce_loss(pred, lab_im)
dice = calc_dice(pred, lab_im)
valid_bce.append(bce.item())
valid_dice.append(dice)
result['valid/BCE'].append(statistics.mean(valid_bce))
result['valid/Dice'].append(statistics.mean(valid_dice))
writer.add_scalar('valid/BinaryCrossEntropy', result['valid/BCE'][-1], epoch+1)
writer.add_scalar('valid/DiceScore', result['valid/Dice'][-1], epoch+1)
print('BCE: {}, Dice: {}'.format(result['valid/BCE'][-1], result['valid/Dice'][-1]))
if best_dice < result['valid/Dice'][-1]:
best_dice = result['valid/Dice'][-1]
best_model_name = os.path.join(args.save_model_path, f'best_model_{epoch + 1:04}.pth')
print('save model ==>> {}'.format(best_model_name))
torch.save(model.state_dict(), best_model_name)
#### create dataloader ####
train_loader = DataLoader(dataset=training_set,
drop_last=True,
sampler=RandomSampler(training_set),
batch_size=batch_size)
valid_loader = DataLoader(dataset=valid_set,
drop_last=False,
sampler=SequentialSampler(valid_set),
batch_size=batch_size)
epochs = 25
model = Unet().cuda()
criterion = DiceBCELoss()
learning_rate = 0.001
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate, weight_decay=1e-5)
train_loss,val_loss = [],[]
train_iou,val_iou = [],[]
valid_loss_min = 10000
i = 1
if initial_checkpoint is not None:
model, optimizer, epochnum, valid_loss_min = load_ckp(checkpoint_path+initial_checkpoint, model, optimizer)
print(f'initial ckp: {epochnum}')
i = i + epochnum
for epoch in range(epochs):
running_train_loss = []
netD_A = Discriminator_MSE(out_ch).cuda(cuda)
netD_B = Discriminator_MSE(in_ch).cuda(cuda)
if lambda_identity > 0.0: # only works when input and output images have the same number of channels
assert (in_ch == out_ch)
fake_A_pool = ImagePool(
pool_size) # create image buffer to store previously generated images
fake_B_pool = ImagePool(
pool_size) # create image buffer to store previously generated images
# define loss functions
criterionGAN = torch.nn.MSELoss(
) # define GAN loss. CrossEntropyLoss or MSELoss !!!!!!!
criterionCycle = torch.nn.L1Loss()
criterionIdt = torch.nn.L1Loss()
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
optimizer_G_A = torch.optim.Adam(netG_A.parameters(),
lr=learning_rate,
betas=(0.5, 0.999))
optimizer_G_B = torch.optim.Adam(netG_B.parameters(),
lr=learning_rate,
betas=(0.5, 0.999))
optimizer_D_A = torch.optim.Adam(netD_A.parameters(),
lr=learning_rate,
betas=(0.5, 0.999))
optimizer_D_B = torch.optim.Adam(netD_B.parameters(),
lr=learning_rate,
betas=(0.5, 0.999))
optimizers = [optimizer_G_A, optimizer_G_B, optimizer_D_A, optimizer_D_B]
realA_mean = np.load("/home/fenqiang/harmonization/realA_mean.npy")
model.eval()
img_list = os.listdir(img_dir)
with torch.no_grad():
for img_name in img_list:
img = cv2.imread(os.path.join(img_dir, img_name))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.transpose((2, 0, 1))
img = np.expand_dims(img, axis=0)
img = np.array(img / 255, dtype=np.float32)
img = torch.from_numpy(img)
img = img.to(torch.device('cuda'))
output = model(img)
output[output >= 0.5] = 1
output[output < 0.5] = 0
output = output.cpu().numpy()
output = np.squeeze(output, 0)
output = np.transpose(output, (1, 2, 0))
output = np.array(output * 255, dtype='uint8')
cv2.imwrite('./result/' + img_name, output)
if __name__ == '__main__':
unet = Unet().to(torch.device('cuda'))
optimizer = optim.Adam(unet.parameters(), lr=0.00001)
train_data = MyDataset(base_path='./data')
train_loader = DataLoader(dataset=train_data, batch_size=1, shuffle=True)
train(unet, train_loader, optimizer, 100)
# model_path = './model/unet.pth'
# img_dir = './test_img'
# inference(unet, model_path, img_dir)
# 네트워크 학습
transform = transforms.Compose([Normalization(mean=0.5, std=0.5), ToTensor()])
dataset_test = Dataset(data_dir=os.path.join(data_dir, 'test'), transform=transform)
loader_test = DataLoader(dataset_test, batch_size=batch_size, shuffle=False, num_workers=8)
# 부수적인 variable
num_data_test = len(dataset_test)
num_batch_test = np.ceil(num_data_test / batch_size)
#네트워크, 파라미터 생성
net = Unet().to(device)
fn_loss = nn.BCEWithLogitsLoss().to(device)
optim = torch.optim.Adam(net.parameters(), lr = lr)
#부수적인 function
fn_tonumpy = lambda x: x.to('cpu').detach().numpy().transpose(0,2,3,1)
fn_denorm = lambda x, mean, std: (x*std) +mean
fn_class = lambda x: 1.0 * (x > 0.5)
# Tensorboard를 위한 SummaryWriter
writer_train = SummaryWriter(log_dir=os.path.join(log_dir, 'train'))
writer_val = SummaryWriter(log_dir=os.path.join(log_dir, 'val'))
#네트워크 저장, 불러오기
def save(ckpt_dir, net, optim, epoch):
if not os.path.exists(ckpt_dir):
os.makedirs(ckpt_dir)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
print("%d/%d,train_loss:%0.3f" %
(step,
(dt_size - 1) // dataload.batch_size + 1, loss.item()))
break
print("epoch %d loss:%0.3f" % (epoch, epoch_loss / step))
return
torch.save(model.state_dict(), 'weights_%d.pth' % epoch)
return model
#训练模型
model = Unet(3, 1).to(device)
batch_size = 1
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters())
dataset = MyDataset(image_path,
label_path,
train_data,
transform=x_transforms,
target_transform=y_transforms)
data_loader = data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=0)
train_model(model, criterion, optimizer, data_loader)
map_location=device)['net'])
model = Unet(n_channels=3, n_classes=8)
model.to(device=device)
#model_dir = './checkpoints/best_score_model_unet.pth'
model_dir = './checkpoints/student_net.pth'
if os.path.exists(model_dir):
#model = load_GPUS(model_dir, model_dir, kwargs)
model.load_state_dict(torch.load(model_dir)['net'])
print("loading model sccessful----" + model_dir)
#model.load_state_dict(torch.load('teach_net_params_0.9895.pkl'))
criterion = nn.CrossEntropyLoss()
criterion2 = nn.KLDivLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0001)
#optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
n_size = batch_size * 256 * 256
writer = SummaryWriter(comment='student' +
f'LR_0.0001_BS_32') #创建一个tensorboard文件
epochs = 50
global_step = 1
for epoch in range(epochs):
loss_sigma = 0.0
correct = 0.0
total = 0.0
model.train()
with tqdm(total=train_steps,
desc=f'Epoch {epoch + 1}/{epochs}',
unit='img') as pbar:
