def train(unet, lr, epochs, batch, CUDA, shuffle, retrain):
if not retrain:
print("Loading model...")
unet.load_state_dict(torch.load("unet.pkl"))
print("OK!")
print("=============\nStart training\nlr: {}\nepoch: {}".format(lr, epochs))
print(
"batch: {}\nshuffle: {}\nCUDA: {}\n=============".format(batch, shuffle, CUDA)
)
optimizer = optim.SGD(unet.parameters(), lr=lr, momentum=0.7)
criterion = nn.BCELoss() # change loss function
tic = time.time()
Tic = tic
for epoch in range(epochs):
epoch_loss = 0
tloader = train_loader(batch, shuffle)
for i, D in enumerate(tloader):
optimizer.zero_grad()
img, msk = D
if CUDA:
img = img.cuda()
msk = msk.cuda()
pred_msk = unet(img).view(-1)
# contiguous的作用是将浅拷贝变成深拷贝，以此使用view
loss = criterion(pred_msk, msk.contiguous().view(-1))
epoch_loss += loss.item()
loss.backward()
optimizer.step()
toc = time.time()
print(
"epoch {} finished --- loss: {:.6f} --- Cost: {:.2f} min".format(
epoch, epoch_loss / i, (toc - tic) / 60
)
)
tic = toc
torch.save(unet.state_dict(), "unet.pkl")
print("Finished. Cost: {:.2f} min".format((tic - Tic) / 60))
if __name__ == "__main__":
parser.add_argument("-v", "--visual", action="store_true")
parser.add_argument("-l", "--lr", type=float, default=1e-5)
parser.add_argument("-e", "--epochs", type=int, default=10)
parser.add_argument("-b", "--batch", type=int, default=40)
parser.add_argument("-r", "--retrain", type=bool, default=False)
args = parser.parse_args()
if args.visual:
from visual import show_pred_mask
from loader import train_loader
trl = train_loader(1, shuffle=True)
img, msk = next(trl)
unet.load_state_dict(torch.load("unet.pkl"))
show_pred_mask(unet, img, msk)
else:
if use_CUDA:
unet.cuda()
try:
train(
unet,
lr=args.lr,
epochs=args.epochs,
batch=args.batch,
CUDA=use_CUDA,
shuffle=True,
def train_thyroid(args):
# construct model
model = None
if args.model_name == "InceptionV3":
model = models.inception_v3(pretrained=args.pretrained)
## Change the last layer
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, args.num_class)
elif args.model_name == "VGG16BN":
model = models.vgg16_bn(pretrained=args.pretrained)
num_ftrs = model.classifier[6].in_features
feature_model = list(model.classifier.children())
feature_model.pop()
feature_model.append(nn.Linear(num_ftrs, args.num_class))
model.classifier = nn.Sequential(*feature_model)
elif args.model_name == "ResNet50":
model = models.resnet50(pretrained=args.pretrained)
in_features = model.fc.in_features
model.fc = torch.nn.Linear(in_features, args.num_class)
else:
raise Exception("Unknown model: {}".format(args.model_name))
# cuda setting
torch.cuda.manual_seed(args.seed)
model.cuda()
cudnn.benchmark = True
# optimizer & loss
# optimizer = optim.SGD(model.parameters(), lr=args.lr,
# weight_decay=5.0e-4, momentum=0.9, nesterov=True)
optimizer = optim.RMSprop(model.parameters(),
lr=args.lr,
alpha=0.99,
weight_decay=1.0e-6,
momentum=0.9)
criterion = nn.CrossEntropyLoss()
# dataloader
train_thyroid_loader = train_loader(args.batch_size)
val_thyroid_loader = val_loader(args.batch_size)
# folder for model saving
if args.pretrained == True:
model_save_dir = os.path.join(args.model_dir, "ft-" + args.model_name,
args.session)
else:
model_save_dir = os.path.join(args.model_dir, "new-" + args.model_name,
args.session)
if os.path.exists(model_save_dir):
shutil.rmtree(model_save_dir)
os.makedirs(model_save_dir)
# viz = Visdom()
# vis_title = "Thyroid classification"
# vis_legend = ['Train Acc', 'Val Acc']
# acc_plot = create_vis_plot(viz, 'Epoch', 'Acc', vis_title, vis_legend)
best_acc = 0.0
for epoch in range(1, args.epochs + 1):
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train_acc = train(train_thyroid_loader, model, criterion, optimizer,
epoch, args)
# evaluate on validation set
val_acc = validate(val_thyroid_loader, model, criterion, args)
is_best = val_acc > best_acc
# update_vis_plot(viz, epoch, train_acc, val_acc, acc_plot, 'append')
if is_best == True:
best_acc = val_acc
cur_model_name = str(epoch).zfill(2) + "-{:.3f}.pth".format(
best_acc)
torch.save(model.cpu(), os.path.join(model_save_dir,
cur_model_name))
print('Save weights at {}/{}'.format(model_save_dir,
cur_model_name))
model.cuda()
loss.backward()
optimizer.step()
if batch_idx > 0 and batch_idx % args.log_interval == 0:
batch_progress = 100. * batch_idx / len(data_loader)
print("Train Epoch: {} [{}/{} ({:.1f})%)]\t Loss: {:.6f}".format(
epoch, batch_idx * len(data), len(data_loader.dataset),
batch_progress, loss.data[0]
))
cur_model_name = args.model_name + "-" + str(epoch).zfill(2) + ".pth"
torch.save(model.state_dict(), os.path.join(args.model_dir, cur_model_name))
print('Save weights at {}/{}'.format(args.model_dir, cur_model_name))
if __name__ == '__main__':
args = set_args()
# Config model and gpu
torch.manual_seed(args.seed)
model = Net()
args.cuda = torch.cuda.is_available()
if args.cuda:
torch.cuda.manual_seed(args.seed)
model.cuda(args.device_id)
import torch.backends.cudnn as cudnn
cudnn.benchmark = True
# dataloader
data_loader = train_loader(args)
# Start training
train(data_loader, model, args)