class Dataset(ABC):
path: str
num_partitions: int
batch_size: int
base_dataset: Optional[TorchDataset]
dataset: TorchDataset
def __init__(
self,
path: str,
num_partitions: int,
batch_size: Optional[int] = None,
base_dataset: Optional[TorchDataset] = None,
) -> None:
if not base_dataset:
raise Exception("Argument 'base_dataset' must be defined")
self.path = path
self.num_partitions = num_partitions
self.batch_size = batch_size or int(128 / num_partitions)
self.partitions = {
i: 1 / num_partitions
for i in range(num_partitions)
}
self.__get_dataset(base_dataset)
def __get_dataset(self, base_dataset: TorchDataset) -> None:
if self.num_partitions > 1:
self.dataset = SplitDataset(base_dataset,
partitions=self.partitions)
else:
self.dataset = base_dataset
def get_train_set(self,
partition: int,
shuffle: bool = True) -> DataLoader:
if self.num_partitions > 1:
self.dataset.select(partition)
return DataLoader(self.dataset,
batch_size=self.batch_size,
shuffle=shuffle)
def main():
parser = get_parser()
args = parser.parse_args()
if not torch.cuda.is_available():
raise ValueError(
"The script requires CUDA support, but CUDA not available")
use_ddp = args.ddp > 0
use_horovod = args.horovod > 0
# Fix seeds in order to get the same losses across runs
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cfg = {
"microbatches": args.num_microbatches,
"placement_strategy": "spread",
"pipeline": args.pipeline,
"optimize": "speed",
"partitions": args.num_partitions,
"horovod": use_horovod,
"ddp": use_ddp,
}
smp.init(cfg)
# SM Distributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")
kwargs = {"batch_size": args.batch_size}
kwargs.update({"num_workers": 1, "pin_memory": True, "shuffle": False})
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])
# SM Distributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
if smp.local_rank() == 0:
dataset1 = datasets.MNIST("../data",
train=True,
download=True,
transform=transform)
smp.barrier()
dataset1 = datasets.MNIST("../data",
train=True,
download=False,
transform=transform)
if (use_ddp or use_horovod) and smp.dp_size() > 1:
partitions_dict = {
f"{i}": 1 / smp.dp_size()
for i in range(smp.dp_size())
}
dataset1 = SplitDataset(dataset1, partitions=partitions_dict)
dataset1.select(f"{smp.dp_rank()}")
# Download and create dataloaders for train and test dataset
dataset2 = datasets.MNIST("../data", train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1, **kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **kwargs)
model = GroupedNet()
# SMP handles the transfer of parameters to the right device
# and the user doesn't need to call 'model.to' explicitly.
# model.to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
# SM Distributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
scaler = smp.amp.GradScaler()
optimizer = smp.DistributedOptimizer(optimizer)
if args.partial_checkpoint:
checkpoint = smp.load(args.partial_checkpoint, partial=True)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
elif args.full_checkpoint:
checkpoint = smp.load(args.full_checkpoint, partial=False)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, scaler, device, train_loader, optimizer, epoch)
test_loss = test(args, model, device, test_loader)
scheduler.step()
if args.save_partial_model:
if smp.dp_rank() == 0:
model_dict = model.local_state_dict()
opt_dict = optimizer.local_state_dict()
smp.save(
{
"model_state_dict": model_dict,
"optimizer_state_dict": opt_dict
},
f"./pt_mnist_checkpoint.pt",
partial=True,
)
if args.save_full_model:
if smp.dp_rank() == 0:
model_dict = model.state_dict()
opt_dict = optimizer.state_dict()
smp.save(
{
"model_state_dict": model_dict,
"optimizer_state_dict": opt_dict
},
"./pt_mnist_checkpoint.pt",
partial=False,
)
# Waiting the save checkpoint to be finished before run another allgather_object
smp.barrier()
if args.assert_losses:
if use_horovod or use_ddp:
# SM Distributed: If using data parallelism, gather all losses across different model
# replicas and check if losses match.
losses = smp.allgather(test_loss, smp.DP_GROUP)
for l in losses:
assert math.isclose(l, losses[0])
assert test_loss < 0.18
else:
assert test_loss < 0.08
def main():
if not torch.cuda.is_available():
raise ValueError(
"The script requires CUDA support, but CUDA not available")
use_ddp = True
use_horovod = False
# Fix seeds in order to get the same losses across runs
random.seed(1)
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
smp.init()
# SM Distributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")
kwargs = {"batch_size": 64}
kwargs.update({"num_workers": 1, "pin_memory": True, "shuffle": False})
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])
# SM Distributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
if smp.local_rank() == 0:
dataset1 = datasets.MNIST("../data",
train=True,
download=True,
transform=transform)
smp.barrier()
dataset1 = datasets.MNIST("../data",
train=True,
download=False,
transform=transform)
if (use_ddp or use_horovod) and smp.dp_size() > 1:
partitions_dict = {
f"{i}": 1 / smp.dp_size()
for i in range(smp.dp_size())
}
dataset1 = SplitDataset(dataset1, partitions=partitions_dict)
dataset1.select(f"{smp.dp_rank()}")
# Download and create dataloaders for train and test dataset
dataset2 = datasets.MNIST("../data", train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1, **kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **kwargs)
model = GroupedNet()
# SMP handles the transfer of parameters to the right device
# and the user doesn't need to call 'model.to' explicitly.
# model.to(device)
optimizer = optim.Adadelta(model.parameters(), lr=4.0)
# SM Distributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
scaler = smp.amp.GradScaler()
optimizer = smp.DistributedOptimizer(optimizer)
scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
for epoch in range(1, 2):
train(model, scaler, device, train_loader, optimizer, epoch)
test_loss = test(model, device, test_loader)
scheduler.step()
if smp.rank() == 0:
if os.path.exists("/opt/ml/local_checkpoints"):
print("-INFO- PATH DO EXIST")
else:
os.makedirs("/opt/ml/local_checkpoints")
print("-INFO- PATH DO NOT EXIST")
# Waiting the save checkpoint to be finished before run another allgather_object
smp.barrier()
if smp.dp_rank() == 0:
model_dict = model.local_state_dict()
opt_dict = optimizer.local_state_dict()
smp.save(
{
"model_state_dict": model_dict,
"optimizer_state_dict": opt_dict
},
f"/opt/ml/local_checkpoints/pt_mnist_checkpoint.pt",
partial=True,
)
smp.barrier()
if smp.local_rank() == 0:
print("Start syncing")
base_s3_path = os.path.dirname(
os.path.dirname(os.getenv("SM_MODULE_DIR", "")))
curr_host = os.getenv("SM_CURRENT_HOST")
full_s3_path = f"{base_s3_path}/checkpoints/{curr_host}/"
sync_local_checkpoints_to_s3(local_path="/opt/ml/local_checkpoints",
s3_path=full_s3_path)
print("Finished syncing")
netG.init_weights()
netD.init_weights()
criterion = nn.BCELoss()
optimizerD = optim.Adam(netD.parameters(),
lr=LEARNING_RATE,
betas=(BETA1, 0.999))
optimizerG = optim.Adam(netG.parameters(),
lr=LEARNING_RATE,
betas=(BETA1, 0.999))
print('Loading dataset...')
images_dataset = ImageDataset(IMAGES_PATH, transforms, REAL_LABEL, device,
'*.*')
real_dataset = SplitDataset(images_dataset, {'train': 0.8, 'validation': 0.2})
real_dataset.select('train')
real_data_loader = DataLoader(real_dataset,
shuffle=True,
batch_size=BATCH_SIZE)
generator_dataset = GeneratorDataset(netG, len(real_dataset), Z_SIZE,
FAKE_LABEL, device)
generator_dataset.generate()
gen_data_loader = DataLoader(generator_dataset,
shuffle=True,
batch_size=BATCH_SIZE)
concat_dataset = ConcatDataset([real_dataset, generator_dataset])
concat_data_loader = DataLoader(concat_dataset,
shuffle=True,
batch_size=BATCH_SIZE)
#real_dataset = dset.LSUN(root='../lsun', classes=['bedroom_train'], transform=transforms)
print('Done loading dataset.')