def __init__(self):
super(MyExperiment, self).__init__()
self.auto_setup()
self.net = UNet(in_channels=1, out_channels=1, dim=3)
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
self.trainloader = trainloader
self.optimizer = optim.Adam(self.net.parameters())
self.criterion = nn.CrossEntropyLoss()
class MyExperiment(BaseExperiment, TensorboardMixin):
def __init__(self):
super(MyExperiment, self).__init__()
self.auto_setup()
self.net = UNet(in_channels=1, out_channels=1, dim=3)
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
self.trainloader = trainloader
self.optimizer = optim.Adam(self.net.parameters())
self.criterion = nn.CrossEntropyLoss()
def run(self):
running_loss = 0.0
for i, data in enumerate(self.trainloader, 0):
# get the inputs
inputs, labels = data
inputs, labels = inputs.to(self.device), labels.to(self.device)
# zero the parameter gradients
self.optimizer.zero_grad()
# forward + backward + optimize
outputs = self.net(inputs)
loss = self.criterion(outputs, labels)
loss.backward()
self.optimizer.step()
# print statistics
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' % (1, i + 1, running_loss / 2000))
running_loss = 0.0
def test_unet_3d(self):
from inferno.extensions.models import UNet
tester = ModelTester((1, 1, 16, 64, 64), (1, 1, 16, 64, 64))
if cuda.is_available():
tester.cuda()
# test default unet 3d
tester(UNet(1, 1, dim=3, initial_features=8))
transforms=trans2, **yaml2dict('config_train.yml')['slicing_config_truth'])
trainset = Zip(imageset_train, labelset_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCHSIZE,
shuffle=True, num_workers=2)
imageset_val = HDF5VolumeLoader(path='./val-volume.h5', path_in_h5_dataset='data',
transforms=trans, **yaml2dict('config_val.yml')['slicing_config'])
labelset_val = HDF5VolumeLoader(path='./stardistance_val.h5', path_in_h5_dataset='data',
transforms=trans2, **yaml2dict('config_val.yml')['slicing_config_truth'])
trainset = Zip(imageset_val, labelset_val)
valloader = torch.utils.data.DataLoader(trainset, batch_size=BATCHSIZE,
shuffle=True, num_workers=2)
net = torch.nn.Sequential(
ConvReLU2D(in_channels=1, out_channels=3, kernel_size=3),
UNet(in_channels=3, out_channels=N_DIRECTIONS, dim=2, final_activation='ReLU')
)
trainer = Trainer(net)
trainer.bind_loader('train', trainloader)
trainer.bind_loader('validate', valloader)
trainer.save_to_directory('./checkpoints')
trainer.save_every((200, 'iterations'))
trainer.build_logger(TensorboardLogger(log_scalars_every=(1, 'iteration'),
log_images_every='never'), log_directory=LOG_DIRECTORY)
trainer.validate_every((200, 'iterations'), for_num_iterations=50)
def test_unet_3d(self):
tester = ModelTester((1, 1, 16, 64, 64), (1, 1, 16, 64, 64))
if cuda.is_available():
tester.cuda()
# test default unet 3d
tester(UNet(1, 1, dim=3, initial_features=8))
def test_unet_2d(self):
tester = ModelTester((1, 1, 256, 256), (1, 1, 256, 256))
if cuda.is_available():
tester.cuda()
tester(UNet(1, 1, dim=2, initial_features=32))
def test_unet_2d(self):
from inferno.extensions.models import UNet
tester = ModelTester((1, 1, 256, 256), (1, 1, 256, 256))
if cuda.is_available():
tester.cuda()
tester(UNet(1, 1, dim=2, initial_features=32))